cancel process

src/api/routes.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+import threading
 
 from fastapi import APIRouter, HTTPException, Request
 from fastapi.responses import StreamingResponse
@@ -25,10 +26,6 @@ async def chat_completions(request: Request):
     except Exception:
         raise HTTPException(status_code=400, detail="Invalid JSON")
 
-    logger.info(
-        f"API Request received. Model: {data.get('model')}, Stream: {data.get('stream', True)}"
-    )
-
     messages = [
         {"role": m.get("role", "user"), "content": get_clean_text(m.get("content"))}
         for m in data.get("messages", [])
@@ -37,19 +34,17 @@ async def chat_completions(request: Request):
     max_tokens = data.get("max_tokens", settings.DEFAULT_MAX_TOKENS)
     temperature = data.get("temperature", settings.DEFAULT_TEMP)
     top_p = data.get("top_p", 0.95)
-
     stop = data.get("stop", [])
     if isinstance(stop, str):
         stop = [stop]
 
-    default_stops = ["<|im_end|>", "<|endoftext|>"]
+    default_stops = ["<|im_end|>", "<|endoftext|>", "<|file_sep|>"]
     for s in default_stops:
         if s not in stop:
             stop.append(s)
 
-    stream_req = data.get("stream", True)
+    abort_event = threading.Event()
 
-    # --- Логика Streaming ---
     async def stream_generator():
         queue = asyncio.Queue()
         loop = asyncio.get_running_loop()
@@ -61,57 +56,60 @@ async def chat_completions(request: Request):
                     "temperature": float(temperature),
                     "top_p": float(top_p),
                     "stop": stop,
+                    "abort_event": abort_event,
                 }
 
+                # Запускаем генерацию
                 for chunk in engine.generate_stream(messages, **gen_kwargs):
                     loop.call_soon_threadsafe(queue.put_nowait, chunk)
-                loop.call_soon_threadsafe(queue.put_nowait, None)  # Конец
+
+                loop.call_soon_threadsafe(queue.put_nowait, None)
             except Exception as e:
-                logger.error(f"Generation error: {e}")
+                if not abort_event.is_set():
+                    logger.error(f"Generation error: {e}")
                 loop.call_soon_threadsafe(queue.put_nowait, {"error": str(e)})
 
         loop.run_in_executor(None, worker)
 
-        while True:
-            chunk = await queue.get()
-
-            if chunk is None:
-                yield "data: [DONE]\n\n"
-                break
-
-            if isinstance(chunk, dict) and "error" in chunk:
-                err_json = json.dumps(
-                    {"error": {"message": chunk["error"], "type": "internal_error"}}
-                )
-                yield f"data: {err_json}\n\n"
-                break
-
-            # Стандартный чанк
-            yield f"data: {json.dumps(chunk)}\n\n"
-
-    if stream_req:
-        headers = {
-            "X-Accel-Buffering": "no",
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-            "Content-Type": "text/event-stream",
-        }
-        return StreamingResponse(
-            stream_generator(), media_type="text/event-stream", headers=headers
-        )
-
-    else:
-
-        def run_sync():
-            with engine.lock:
-                return engine.llm.create_chat_completion(
-                    messages=messages,
-                    max_tokens=int(max_tokens),
-                    temperature=float(temperature),
-                    top_p=float(top_p),
-                    stop=stop,
-                    stream=False,
-                )
-
-        response = await asyncio.to_thread(run_sync)
-        return response
+        try:
+            while True:
+                if await request.is_disconnected():
+                    logger.info("Client disconnected! Aborting generation...")
+                    abort_event.set()
+                    break
+
+                try:
+                    chunk = await asyncio.wait_for(queue.get(), timeout=0.1)
+                except asyncio.TimeoutError:
+                    continue
+
+                if chunk is None:
+                    yield "data: [DONE]\n\n"
+                    break
+
+                if isinstance(chunk, dict) and "error" in chunk:
+                    if abort_event.is_set():
+                        break
+                    err_json = json.dumps(
+                        {"error": {"message": chunk["error"], "type": "internal_error"}}
+                    )
+                    yield f"data: {err_json}\n\n"
+                    break
+
+                yield f"data: {json.dumps(chunk)}\n\n"
+
+        except asyncio.CancelledError:
+            logger.info("Task cancelled. Stopping worker.")
+            abort_event.set()
+            raise
+
+    # Возвращаем стрим
+    headers = {
+        "X-Accel-Buffering": "no",
+        "Cache-Control": "no-cache",
+        "Connection": "keep-alive",
+        "Content-Type": "text/event-stream",
+    }
+    return StreamingResponse(
+        stream_generator(), media_type="text/event-stream", headers=headers
+    )

src/core/engine.py

@@ -1,5 +1,6 @@
 import threading
-from typing import Any, Dict, Generator, List
+import time
+from typing import Any, Dict, Generator, List, Optional
 
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
@@ -24,14 +25,18 @@ class ModelEngine:
                 n_ctx=settings.CONTEXT_SIZE,
                 n_threads=settings.N_THREADS,
                 n_gpu_layers=settings.N_GPU_LAYERS,
-                verbose=False,
+                verbose=True,
             )
             print("Model loaded successfully!")
         except Exception as e:
             print(f"CRITICAL ERROR loading model: {e}")
 
-    # Изменили сигнатуру: теперь принимает **kwargs
-    def generate_stream(self, messages: List[Dict[str, str]], **kwargs) -> Generator:
+    def generate_stream(
+        self,
+        messages: List[Dict[str, str]],
+        abort_event: Optional[threading.Event] = None,  # Новый аргумент
+        **kwargs,
+    ) -> Generator:
         if not self.llm:
             raise RuntimeError("Model not loaded")
 
@@ -39,7 +44,18 @@ class ModelEngine:
         temperature = kwargs.get("temperature", settings.DEFAULT_TEMP)
         stop = kwargs.get("stop", [])
 
-        with self.lock:
+        acquired = False
+        while not acquired:
+            if abort_event and abort_event.is_set():
+                print("Request aborted while waiting in queue.")
+                return
+
+            acquired = self.lock.acquire(timeout=0.5)
+
+        try:
+            if abort_event and abort_event.is_set():
+                return
+
             stream = self.llm.create_chat_completion(
                 messages=messages,
                 max_tokens=int(max_tokens),
@@ -48,8 +64,16 @@ class ModelEngine:
                 stream=True,
                 top_p=kwargs.get("top_p", 0.95),
             )
+
             for chunk in stream:
+                if abort_event and abort_event.is_set():
+                    print("Request aborted during generation.")
+                    break
+
                 yield chunk
 
+        finally:
+            self.lock.release()
+
 
 engine = ModelEngine()