feat(core): implement thread-safe generation abort mechanism

- Add `AbortCriteria` class and a thread-safe `Llama.abort()` method to allow graceful interruption of ongoing text generation from external threads (e.g., UI or async environments).
- Automatically inject `AbortCriteria` into the stopping criteria sequence at the start of `_create_completion`.
- Ensure that when an abort is triggered, the partially generated `completion_tokens` are correctly detokenized and preserved.
- Set `finish_reason` to `"abort"` when generation is interrupted, allowing downstream streaming clients to correctly identify manual cancellations.
- Simplify and optimize the stopping criteria evaluation logic within the core `generate` loop.
- Reorganize and sort module imports for better readability.

Signed-off-by: JamePeng <jame_peng@sina.com>

Author: JamePeng

llama_cpp/llama.py

@@ -1,17 +1,20 @@
 from __future__ import annotations
 
+import contextlib
+import ctypes
+import fnmatch
+import json
+import multiprocessing
 import os
 import sys
-import uuid
 import time
-import json
-import ctypes
+import threading
 import typing
-import random
-import fnmatch
+import uuid
 import warnings
-import contextlib
-import multiprocessing
+
+import numpy as np
+import numpy.typing as npt
 
 from typing import (
     Any,
@@ -29,7 +32,6 @@
 from collections import deque
 from pathlib import Path
 
-
 from .llama_types import *
 from .llama_grammar import LlamaGrammar
 from .llama_cache import (
@@ -46,9 +48,6 @@
 
 from llama_cpp.llama_speculative import LlamaDraftModel
 
-import numpy as np
-import numpy.typing as npt
-
 import llama_cpp._internals as internals
 from ._internals import (
     LlamaSamplingContext,
@@ -60,6 +59,19 @@
 from ._utils import suppress_stdout_stderr
 
 
+class AbortCriteria:
+    """
+    Listen for external interruption signals to trigger a stop condition.
+    When an external thread calls `llama.abort()`, a loop interrupt is generated.
+    """
+    def __init__(self, abort_event: threading.Event):
+        self.abort_event = abort_event
+
+    def __call__(self, _input_ids: npt.NDArray[np.intc], _logits: npt.NDArray[np.single]) -> bool:
+        # Note: _input_ids and _logits are required by the signature but unused here.
+        return self.abort_event.is_set()
+
+
 class Llama:
     """High-level Python wrapper for a llama.cpp model."""
 
@@ -623,6 +635,9 @@ def __init__(
 
         self._sampling_ctx: Optional[LlamaSamplingContext] = None
 
+        # Create a thread-safe interrupt event
+        self._abort_event = threading.Event()
+
     def close(self) -> None:
         """Explicitly free the model from memory."""
         if getattr(self, "_sampling_ctx", None) is not None:
@@ -769,6 +784,15 @@ def reset(self):
         """Reset the model state."""
         self.n_tokens = 0
 
+    def abort(self) -> None:
+        """
+        Safely aborts any ongoing text generation.
+        Useful for async API environments or UI interruption buttons.
+        """
+        if self.verbose:
+            print(f"Llama.abort: Abort signal received. Terminating generation...", file=sys.stderr)
+        self._abort_event.set()
+
     def eval(
             self,
             tokens: Sequence[int],
@@ -1442,26 +1466,18 @@ def adapter(token_data_array: llama_cpp.llama_token_data_array):
 
                 # Sample loop
                 while sample_idx < self.n_tokens:
+                    if self._abort_event.is_set():
+                        return
+
                     token = self._sampling_ctx.sample(self._ctx, idx=-1)
                     self._sampling_ctx.accept(token, False if grammar is None else True)
 
                     sample_idx += 1
 
-                    # Halt generation if custom stopping criteria are met
                     if stopping_criteria is not None:
-                        if self._logits_all:
-                            logits_idx = sample_idx - self.n_tokens
-                            check_stopping = True
-                        else:
-                            if sample_idx == self.n_tokens:
-                                logits_idx = 0
-                                check_stopping = True
-                            else:
-                                check_stopping = False
-
-                        if check_stopping and stopping_criteria(
+                        if stopping_criteria(
                             self._input_ids[: sample_idx],
-                            self._scores[logits_idx, :]
+                            self._scores[0 if not self._logits_all else sample_idx - self.n_tokens, :]
                         ):
                             return
 
@@ -1746,6 +1762,8 @@ def _create_completion(
         Iterator[CreateCompletionResponse], Iterator[CreateCompletionStreamResponse]
     ]:
         assert suffix is None or suffix.__class__ is str
+        # Each time a new request is initiated, the previous abort state must be cleared.
+        self._abort_event.clear()
 
         completion_id: str = f"cmpl-{str(uuid.uuid4())}"
         created: int = int(time.time())
@@ -1885,6 +1903,11 @@ def _create_completion(
                 if self.verbose:
                     print("Llama._create_completion: cache miss", file=sys.stderr)
 
+        if stopping_criteria is None:
+            stopping_criteria = StoppingCriteriaList([AbortCriteria(self._abort_event)])
+        else:
+            stopping_criteria.append(AbortCriteria(self._abort_event))
+
         finish_reason = "length"
         multibyte_fix = 0
         for token in self.generate(
@@ -1929,6 +1952,11 @@ def _create_completion(
                 finish_reason = "stop"
                 break
 
+            if self._abort_event.is_set():
+                text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
+                finish_reason = "abort"
+                break
+
             completion_tokens.append(token)
 
             all_text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
@@ -2108,6 +2136,11 @@ def _create_completion(
             text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
             finish_reason = "stop"
 
+        # If the abort is triggered externally, force the `finish_reason` to be changed to "abort".
+        if self._abort_event.is_set():
+            text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
+            finish_reason = "abort"
+
         if self.verbose:
             self._ctx.print_timings()