fix ICL pool rules for version and exemple counter

app.py

@@ -5,6 +5,7 @@ file exposes a small REST API and a tiny in-memory session store. State is
 ephemeral and per-process; perfect for a single-user demo or HF Space.
 """
 from __future__ import annotations
+from copy import deepcopy
 
 import asyncio
 import os
@@ -362,22 +363,68 @@ def reset_all():
 
 
 # --- token edit ------------------------------------------------------------
+def _add_or_update_sentence_in_icl(idx: int) -> str:
+    sents = SESSION["sentences"]
+    if idx < 0 or idx >= len(sents):
+        raise HTTPException(404, "Bad sentence idx")
+    sent = sents[idx]
+    schema_obj = schema_from_dict(SESSION["schema"])
+    pool: ICLPool = SESSION["icl_pool"]
+    tokens_snapshot = deepcopy(sent["tokens"])
+    ann = {
+        "sentence_id": sent["id"],
+        "language": sent["language"] or SESSION["language"],
+        "tokens": tokens_snapshot,
+    }
+
+    result = pool.add(ICLExample(
+        language=sent["language"] or SESSION["language"] or "",
+        schema_hash=schema_obj.hash(),
+        tokens=[t["surface"] for t in tokens_snapshot],
+        gold_annotation=ann,
+        source="corrected",
+    ))
+
+    sent["validated"] = True
+    return result
+
 
 @app.post("/api/sentence/{idx}/token/{tidx}")
 def update_token(idx: int, tidx: int, req: TokenUpdateReq):
     sents = SESSION["sentences"]
+
     if idx < 0 or idx >= len(sents):
         raise HTTPException(404, "Bad sentence idx")
     if tidx < 0 or tidx >= len(sents[idx]["tokens"]):
         raise HTTPException(404, "Bad token idx")
-    # Preserve surface (never editable)
-    surface = sents[idx]["tokens"][tidx]["surface"]
+
+    sent = sents[idx]
+    was_validated = bool(sent.get("validated"))
+
+    surface = sent["tokens"][tidx]["surface"]
     new_tok = {**req.token, "surface": surface}
-    sents[idx]["tokens"][tidx] = new_tok
-    # Remove this token from disagreement list if it was there
-    sents[idx]["disagreements"] = [d for d in sents[idx]["disagreements"] if d["token_idx"] != tidx]
-    sents[idx]["n_disagreements"] = len(sents[idx]["disagreements"])
-    return sents[idx]
+    sent["tokens"][tidx] = new_tok
+
+    sent["disagreements"] = [
+        d for d in sent["disagreements"]
+        if d["token_idx"] != tidx
+    ]
+    sent["n_disagreements"] = len(sent["disagreements"])
+
+    icl_result = None
+
+    # If sentence in ICL pool already, update it. If not, add it. This way we keep the pool in sync with user corrections.
+    if was_validated:
+        icl_result = _add_or_update_sentence_in_icl(idx)
+
+    state = _public_state()
+    state["updated_sentence_idx"] = idx
+    state["icl_add_result"] = icl_result
+    state["icl_duplicate"] = icl_result == "unchanged"
+    state["icl_updated"] = icl_result == "updated"
+    state["icl_inserted"] = icl_result == "inserted"
+
+    return state
 
 
 @app.post("/api/bulk_similar")
@@ -449,30 +496,16 @@ def bulk_update(idx: int, payload: dict):
 
 
 # --- ICL pool --------------------------------------------------------------
-
 @app.post("/api/sentence/{idx}/add_to_icl")
 def add_sentence_to_icl(idx: int):
-    sents = SESSION["sentences"]
-    if idx < 0 or idx >= len(sents):
-        raise HTTPException(404, "Bad sentence idx")
-    sent = sents[idx]
-    schema_obj = schema_from_dict(SESSION["schema"])
-    pool: ICLPool = SESSION["icl_pool"]
-    ann = {
-        "sentence_id": sent["id"],
-        "language": sent["language"] or SESSION["language"],
-        "tokens": sent["tokens"],
-    }
-    pool.add(ICLExample(
-        language=sent["language"] or SESSION["language"] or "",
-        schema_hash=schema_obj.hash(),
-        tokens=[t["surface"] for t in sent["tokens"]],
-        gold_annotation=ann,
-        source="corrected",
-    ))
-    # Adding to ICL implies the user accepts this annotation as gold → mark validated.
-    sent["validated"] = True
-    return _public_state()
+    result = _add_or_update_sentence_in_icl(idx)
+
+    state = _public_state()
+    state["icl_add_result"] = result
+    state["icl_duplicate"] = result == "unchanged"
+    state["icl_updated"] = result == "updated"
+    state["icl_inserted"] = result == "inserted"
+    return state
 
 
 @app.post("/api/sentence/{idx}/sent_score")

prompts.py

@@ -5,6 +5,7 @@ written material. ICLPool keeps a session-scoped, filterable bank of validated
 or corrected examples.
 """
 from __future__ import annotations
+from copy import deepcopy
 
 import json
 import random
@@ -30,6 +31,8 @@ class ICLExample:
     note: str = ""
 
 
+
+
 @dataclass
 class ICLPool:
     """Session-scoped pool of in-context examples.
@@ -40,9 +43,30 @@ class ICLPool:
     entries: list[ICLExample] = field(default_factory=list)
     version: int = 0
 
-    def add(self, ex: ICLExample) -> None:
+    def _key(self, ex: ICLExample) -> tuple[str, str, tuple[str, ...]]:
+        return (
+            ex.language or "",
+            ex.schema_hash or "",
+            tuple(ex.tokens or []),
+        )
+
+    def _same_content(self, a: ICLExample, b: ICLExample) -> bool:
+        return a.gold_annotation == b.gold_annotation
+
+    def add(self, ex: ICLExample) -> str:
+        ex = deepcopy(ex)
+        key = self._key(ex)
+
+        for i, existing in enumerate(self.entries):
+            if self._key(existing) == key:
+                if self._same_content(existing, ex):
+                    return "unchanged"
+                self.entries[i] = ex
+                self.version += 1
+                return "updated"
         self.entries.append(ex)
         self.version += 1
+        return "inserted"
 
     def filter(self, language: str = "", schema_hash: str = "") -> list[ICLExample]:
         out = self.entries

provider.py

@@ -130,6 +130,7 @@ class LLMClient:
             timeout: float = DEFAULT_TIMEOUT,
     ) -> ModelResult:
         """Call one model, validate JSON. One retry on schema-validation failure."""
+        print(f"[LLM] start provider={self.provider} model={model}")
         json_schema = to_json_schema(schema)
         start = time.time()
         msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
@@ -149,13 +150,17 @@ class LLMClient:
                     raw_text = await self._call(client, msgs, json_schema, model, temperature)
                     ann, err = self._parse_and_validate(raw_text, schema)
                 if err:
+                    print(
+                        f"[LLM] error provider={self.provider} model={model} latency={time.time() - start:.2f}s error={e}")
                     return ModelResult(model=model, ok=False, annotation=None, latency_s=time.time() - start, error=err,
                                        raw=raw_text)
+                print(f"[LLM] done provider={self.provider} model={model} latency={time.time() - start:.2f}s")
                 return ModelResult(model=model, ok=True, annotation=ann, latency_s=time.time() - start, raw=raw_text)
             finally:
                 if close_after:
                     await client.aclose()
         except Exception as e:
+            print(f"[LLM] error provider={self.provider} model={model} latency={time.time() - start:.2f}s error={e}")
             return ModelResult(model=model, ok=False, annotation=None, latency_s=time.time() - start, error=str(e))
 
     async def annotate_many(
@@ -180,17 +185,25 @@ class LLMClient:
                     temperature: float) -> str:
         # Strict json_schema works on OpenAI and most OpenRouter models. For Mistral and
         # for some open-source models routed via OpenRouter, fall back to json_object.
-        if self.provider == "mistral":
+        if self.provider in {"mistral", "ilaas"}:
             payload = {
-                "model": model, "messages": msgs, "temperature": temperature,
+                "model": model,
+                "messages": msgs,
+                "temperature": temperature,
                 "response_format": {"type": "json_object"},
             }
         else:
             payload = {
-                "model": model, "messages": msgs, "temperature": temperature,
+                "model": model,
+                "messages": msgs,
+                "temperature": temperature,
                 "response_format": {
                     "type": "json_schema",
-                    "json_schema": {"name": "annotation", "strict": True, "schema": json_schema},
+                    "json_schema": {
+                        "name": "annotation",
+                        "strict": True,
+                        "schema": json_schema,
+                    },
                 },
             }
         resp = await client.post(self.endpoint, headers=self.headers, json=payload)

static/app.js

@@ -769,8 +769,31 @@ function annotator() {
 
         async addSentenceToIcl(sidx) {
             const r = await fetch(`/api/sentence/${sidx}/add_to_icl`, {method: 'POST'});
-            this.applyState(await r.json());
-            this.toast(`Added to ICL pool (v${this.state.icl_pool.version}, ${this.state.icl_pool.size} entries).`, 'ok');
+
+            if (!r.ok) {
+                this.toast('Could not add to ICL pool.', 'error');
+                return;
+            }
+
+            const data = await r.json();
+            this.applyState(data);
+
+            if (data.icl_add_result === 'unchanged') {
+                this.toast(
+                    `Already in ICL pool — unchanged (v${this.state.icl_pool.version}, ${this.state.icl_pool.size} entries).`,
+                    'warn'
+                );
+            } else if (data.icl_add_result === 'updated') {
+                this.toast(
+                    `Updated existing ICL example after correction (v${this.state.icl_pool.version}, ${this.state.icl_pool.size} entries).`,
+                    'ok'
+                );
+            } else {
+                this.toast(
+                    `Added to ICL pool (v${this.state.icl_pool.version}, ${this.state.icl_pool.size} entries).`,
+                    'ok'
+                );
+            }
         },
 
         async setValidated(sidx, value) {
@@ -889,60 +912,97 @@ function annotator() {
         },
 
         async saveToken() {
-            const sidx = this.editor.sidx, tidx = this.editor.tidx;
-            const surface = this.editor.tok.surface;
-            const changes = this.fieldChanges();
-            const wantPropagate = this.editor.propagateToSimilar && Object.keys(changes).length > 0 && this.matchingTokenCount() > 0;
-
-            this.editor.tok._corrected = true;
-            const r = await fetch(`/api/sentence/${sidx}/token/${tidx}`, {
+    const sidx = this.editor.sidx;
+    const tidx = this.editor.tidx;
+    const surface = this.editor.tok.surface;
+    const changes = this.fieldChanges();
+    const wantPropagate =
+        this.editor.propagateToSimilar &&
+        Object.keys(changes).length > 0 &&
+        this.matchingTokenCount() > 0;
+
+    this.editor.tok._corrected = true;
+
+    const r = await fetch(`/api/sentence/${sidx}/token/${tidx}`, {
+        method: 'POST',
+        headers: {'Content-Type': 'application/json'},
+        body: JSON.stringify({token: this.editor.tok})
+    });
+
+    if (!r.ok) {
+        this.toast('Save failed.', 'error');
+        return;
+    }
+
+    // returns full state to ensure consistency
+    const data = await r.json();
+    this.applyState(data);
+
+    let propagatedCount = 0;
+
+    if (wantPropagate) {
+        try {
+            const r2 = await fetch('/api/bulk_similar', {
                 method: 'POST',
                 headers: {'Content-Type': 'application/json'},
-                body: JSON.stringify({token: this.editor.tok})
+                body: JSON.stringify({
+                    surface,
+                    updates: changes,
+                    exclude: [{s: sidx, t: tidx}],
+                }),
             });
-            if (!r.ok) {
-                this.toast('Save failed.', 'error');
-                return;
-            }
-            const sent = await r.json();
-            this.replaceSentence(sidx, sent);
 
-            let propagatedCount = 0;
-            if (wantPropagate) {
-                try {
-                    const r2 = await fetch('/api/bulk_similar', {
-                        method: 'POST',
-                        headers: {'Content-Type': 'application/json'},
-                        body: JSON.stringify({
-                            surface,
-                            updates: changes,
-                            exclude: [{s: sidx, t: tidx}],
-                        }),
-                    });
-                    if (r2.ok) {
-                        const j = await r2.json();
-                        for (const item of (j.sentences || [])) {
-                            this.replaceSentence(item.idx, item.sentence);
-                        }
-                        propagatedCount = (j.affected || []).length;
-                    }
-                } catch (e) {
-                    this.toast('Propagation failed: ' + e.message, 'error');
-                }
-            }
+            if (r2.ok) {
+                const j = await r2.json();
 
-            // auto-advance
-            if (this.editor.autoAdvance) {
-                const next = this.findNextDisagreement(sidx, tidx);
-                if (next) {
-                    this.openTokenEditor(next.s, next.t);
-                    if (propagatedCount > 0) this.toast(`✓ Saved + propagated to ${propagatedCount} other "${surface}".`, 'ok');
-                    return;
+                for (const item of (j.sentences || [])) {
+                    this.replaceSentence(item.idx, item.sentence);
                 }
+
+                propagatedCount = (j.affected || []).length;
             }
-            this.closeModal();
-            this.toast(propagatedCount > 0 ? `✓ Saved + propagated to ${propagatedCount} other "${surface}".` : '✓ Saved.', 'ok');
-        },
+        } catch (e) {
+            this.toast('Propagation failed: ' + e.message, 'error');
+        }
+    }
+
+    let iclMsg = '';
+
+    if (data.icl_add_result === 'updated') {
+        iclMsg = ` + updated ICL v${this.state.icl_pool.version}`;
+    } else if (data.icl_add_result === 'inserted') {
+        iclMsg = ` + added to ICL v${this.state.icl_pool.version}`;
+    } else if (data.icl_add_result === 'unchanged') {
+        iclMsg = ` + ICL unchanged`;
+    }
+
+    // auto-advance
+    if (this.editor.autoAdvance) {
+        const next = this.findNextDisagreement(sidx, tidx);
+
+        if (next) {
+            this.openTokenEditor(next.s, next.t);
+
+            this.toast(
+                propagatedCount > 0
+                    ? `✓ Saved + propagated to ${propagatedCount} other "${surface}"${iclMsg}.`
+                    : `✓ Saved${iclMsg}.`,
+                'ok'
+            );
+
+            return;
+        }
+    }
+
+    this.closeModal();
+
+    this.toast(
+        propagatedCount > 0
+            ? `✓ Saved + propagated to ${propagatedCount} other "${surface}"${iclMsg}.`
+            : `✓ Saved${iclMsg}.`,
+        'ok'
+    );
+},
 
         findNextDisagreement(sidx, tidx) {
             const sents = this.state.sentences;