| |
| """ |
| BlitzKode backend server. |
| |
| Serves the bundled frontend and proxies prompts to a local GGUF model |
| through llama.cpp. Model is loaded lazily so the module stays importable |
| in tests and environments where the model artifact is not present yet. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import asyncio |
| import json |
| import logging |
| import os |
| import time |
| from concurrent.futures import ThreadPoolExecutor |
| from contextlib import asynccontextmanager |
| from dataclasses import dataclass |
| from pathlib import Path |
| from threading import Lock |
| from typing import Iterator |
|
|
| import llama_cpp |
| import uvicorn |
| from fastapi import FastAPI, HTTPException, Request |
| from fastapi.middleware.cors import CORSMiddleware |
| from fastapi.responses import FileResponse, JSONResponse, StreamingResponse |
| from pydantic import BaseModel, Field |
|
|
| APP_NAME = "BlitzKode" |
| APP_VERSION = "2.0" |
| CREATOR = "Sajad" |
| ROOT_DIR = Path(__file__).resolve().parent |
| DEFAULT_MODEL_PATH = ROOT_DIR / "blitzkode.gguf" |
| DEFAULT_FRONTEND_PATH = ROOT_DIR / "frontend" / "index.html" |
| DEFAULT_CONTEXT = 2048 |
| DEFAULT_MAX_PROMPT_LENGTH = 4000 |
| DEFAULT_MAX_TOKENS = 512 |
| STOP_TOKENS = ["<|im_end|>", "<|im_start|>user"] |
|
|
| SYSTEM_PROMPT = ( |
| "<|im_start|>system\n" |
| "You are BlitzKode, an AI coding assistant created by Sajad. " |
| "You are an expert in Python, JavaScript, Java, C++, and other programming languages. " |
| "Write clean, efficient, and well-documented code. Keep responses concise and practical.<|im_end|>" |
| ) |
|
|
| logger = logging.getLogger("blitzkode") |
|
|
|
|
| def _bool_from_env(name: str, default: bool = False) -> bool: |
| value = os.getenv(name) |
| if value is None: |
| return default |
| return value.strip().lower() in {"1", "true", "yes", "on"} |
|
|
|
|
| def _int_from_env(name: str, default: int) -> int: |
| value = os.getenv(name) |
| if not value: |
| return default |
| try: |
| return int(value) |
| except ValueError: |
| return default |
|
|
|
|
| def _validate_prompt(prompt: str, max_length: int) -> tuple[str, JSONResponse | None]: |
| prompt = prompt.strip() |
| if not prompt: |
| return prompt, JSONResponse({"error": "Prompt is required"}, status_code=400) |
| if len(prompt) > max_length: |
| return prompt, JSONResponse( |
| {"error": f"Prompt too long. Max {max_length} chars."}, |
| status_code=400, |
| ) |
| return prompt, None |
|
|
|
|
| @dataclass(slots=True) |
| class Settings: |
| root_dir: Path = ROOT_DIR |
| model_path: Path = Path(os.getenv("BLITZKODE_MODEL_PATH", DEFAULT_MODEL_PATH)) |
| frontend_path: Path = Path(os.getenv("BLITZKODE_FRONTEND_PATH", DEFAULT_FRONTEND_PATH)) |
| host: str = os.getenv("BLITZKODE_HOST", "0.0.0.0") |
| port: int = _int_from_env("BLITZKODE_PORT", 7860) |
| n_gpu_layers: int = _int_from_env("BLITZKODE_GPU_LAYERS", 0) |
| n_ctx: int = _int_from_env("BLITZKODE_N_CTX", DEFAULT_CONTEXT) |
| n_threads: int = _int_from_env("BLITZKODE_THREADS", max(1, min(8, os.cpu_count() or 1))) |
| n_batch: int = _int_from_env("BLITZKODE_BATCH", 128) |
| max_prompt_length: int = _int_from_env("BLITZKODE_MAX_PROMPT_LENGTH", DEFAULT_MAX_PROMPT_LENGTH) |
| preload_model: bool = _bool_from_env("BLITZKODE_PRELOAD_MODEL", default=False) |
| workers: int = _int_from_env("BLITZKODE_WORKERS", 2) |
| cors_origins: str = os.getenv("BLITZKODE_CORS_ORIGINS", "*") |
| api_key: str = os.getenv("BLITZKODE_API_KEY", "") |
|
|
|
|
| class MessageItem(BaseModel): |
| role: str |
| content: str |
|
|
|
|
| class GenerateRequest(BaseModel): |
| prompt: str |
| messages: list[MessageItem] = Field(default_factory=list) |
| temperature: float = Field(default=0.5, ge=0.0, le=2.0) |
| max_tokens: int = Field(default=256, ge=1, le=DEFAULT_MAX_TOKENS) |
| top_p: float = Field(default=0.95, gt=0.0, le=1.0) |
| top_k: int = Field(default=20, ge=1, le=200) |
| repeat_penalty: float = Field(default=1.05, ge=0.8, le=2.0) |
|
|
|
|
| class ModelService: |
| def __init__(self, settings: Settings): |
| self.settings = settings |
| self._llm = None |
| self._lock = Lock() |
| self._load_time_seconds: float | None = None |
| self._last_error: str | None = None |
|
|
| @property |
| def model_loaded(self) -> bool: |
| return self._llm is not None |
|
|
| @property |
| def model_exists(self) -> bool: |
| return self.settings.model_path.exists() |
|
|
| @property |
| def last_error(self) -> str | None: |
| return self._last_error |
|
|
| @property |
| def load_time_seconds(self) -> float | None: |
| return self._load_time_seconds |
|
|
| def load_model(self): |
| if self._llm is not None: |
| return self._llm |
|
|
| with self._lock: |
| if self._llm is not None: |
| return self._llm |
|
|
| if not self.model_exists: |
| self._last_error = f"Model not found at {self.settings.model_path}" |
| raise FileNotFoundError(self._last_error) |
|
|
| start_time = time.perf_counter() |
| try: |
| self._llm = llama_cpp.Llama( |
| model_path=str(self.settings.model_path), |
| n_gpu_layers=self.settings.n_gpu_layers, |
| n_ctx=self.settings.n_ctx, |
| n_threads=self.settings.n_threads, |
| n_batch=self.settings.n_batch, |
| verbose=False, |
| use_mmap=True, |
| use_mlock=False, |
| seed=-1, |
| ) |
| self._load_time_seconds = time.perf_counter() - start_time |
| self._last_error = None |
| logger.info("Model loaded in %.2fs (gpu_layers=%d)", self._load_time_seconds, self.settings.n_gpu_layers) |
| except Exception as exc: |
| self._last_error = str(exc) |
| logger.error("Model load failed: %s", exc) |
| raise |
|
|
| return self._llm |
|
|
| def build_prompt(self, req: GenerateRequest) -> str: |
| parts = [SYSTEM_PROMPT] |
| for msg in req.messages: |
| if msg.role in ("user", "assistant"): |
| parts.append(f"<|im_start|>{msg.role}\n{msg.content}<|im_end|>") |
| parts.append(f"<|im_start|>user\n{req.prompt}<|im_end|>") |
| parts.append("<|im_start|>assistant\n") |
| return "\n".join(parts) |
|
|
| def _gen_params(self, req: GenerateRequest) -> dict: |
| return dict( |
| max_tokens=req.max_tokens, |
| temperature=req.temperature, |
| top_p=req.top_p, |
| top_k=req.top_k, |
| repeat_penalty=req.repeat_penalty, |
| frequency_penalty=0.0, |
| presence_penalty=0.0, |
| stop=STOP_TOKENS, |
| ) |
|
|
| def generate_once(self, req: GenerateRequest) -> dict[str, object]: |
| llm = self.load_model() |
| start = time.perf_counter() |
|
|
| result = llm(self.build_prompt(req), **self._gen_params(req)) |
| response = result["choices"][0]["text"].strip() |
| elapsed = time.perf_counter() - start |
| logger.info("Generated %d chars in %.2fs", len(response), elapsed) |
|
|
| return {"response": response, "creator": CREATOR, "model": APP_NAME, "version": APP_VERSION} |
|
|
| def stream_tokens(self, req: GenerateRequest) -> Iterator[str]: |
| llm = self.load_model() |
| start = time.perf_counter() |
| token_count = 0 |
|
|
| try: |
| for token in llm(self.build_prompt(req), stream=True, **self._gen_params(req)): |
| if not token.get("choices"): |
| continue |
| text = token["choices"][0].get("text", "") |
| if text: |
| token_count += 1 |
| yield f"data: {json.dumps({'token': text})}\n\n" |
| elapsed = time.perf_counter() - start |
| logger.info("Streamed %d tokens in %.2fs", token_count, elapsed) |
| yield "data: [DONE]\n\n" |
| except Exception as exc: |
| logger.error("Stream error: %s", exc) |
| yield f"data: {json.dumps({'error': str(exc)})}\n\n" |
|
|
|
|
| def _check_api_key(request: Request, settings: Settings) -> JSONResponse | None: |
| if not settings.api_key: |
| return None |
| auth = request.headers.get("Authorization", "") |
| token = auth[7:] if auth.startswith("Bearer ") else auth |
| if token != settings.api_key: |
| return JSONResponse({"error": "Unauthorized"}, status_code=401) |
| return None |
|
|
|
|
| def create_app(settings: Settings | None = None) -> FastAPI: |
| settings = settings or Settings() |
| model_service = ModelService(settings) |
| executor = ThreadPoolExecutor(max_workers=settings.workers) |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(message)s", datefmt="%H:%M:%S") |
|
|
| @asynccontextmanager |
| async def lifespan(_: FastAPI): |
| if settings.preload_model: |
| try: |
| await asyncio.to_thread(model_service.load_model) |
| except Exception: |
| pass |
| try: |
| yield |
| finally: |
| executor.shutdown(wait=False, cancel_futures=True) |
|
|
| app = FastAPI(title=f"{APP_NAME} API", version=APP_VERSION, lifespan=lifespan) |
| app.state.settings = settings |
| app.state.model_service = model_service |
| app.state.executor = executor |
|
|
| cors_origins = [o.strip() for o in settings.cors_origins.split(",") if o.strip()] |
| app.add_middleware(CORSMiddleware, allow_origins=cors_origins, allow_methods=["*"], allow_headers=["*"]) |
|
|
| @app.get("/") |
| async def root(): |
| if not settings.frontend_path.exists(): |
| raise HTTPException(status_code=404, detail="Frontend file is missing.") |
| return FileResponse(str(settings.frontend_path)) |
|
|
| @app.get("/health") |
| async def health(): |
| status = "healthy" |
| if not settings.frontend_path.exists() or not model_service.model_exists: |
| status = "degraded" |
| return JSONResponse({ |
| "status": status, |
| "model_loaded": model_service.model_loaded, |
| "model_path": str(settings.model_path), |
| "model_exists": model_service.model_exists, |
| "frontend_exists": settings.frontend_path.exists(), |
| "version": APP_VERSION, |
| "gpu_layers": settings.n_gpu_layers, |
| "last_error": model_service.last_error, |
| }) |
|
|
| @app.post("/generate") |
| async def generate(req: GenerateRequest, request: Request): |
| auth_err = _check_api_key(request, settings) |
| if auth_err: |
| return auth_err |
|
|
| prompt, err = _validate_prompt(req.prompt, settings.max_prompt_length) |
| if err: |
| return err |
|
|
| try: |
| sanitized = req.model_copy(update={"prompt": prompt}) |
| payload = await asyncio.get_running_loop().run_in_executor(executor, model_service.generate_once, sanitized) |
| return JSONResponse(payload) |
| except FileNotFoundError as exc: |
| return JSONResponse({"error": str(exc)}, status_code=503) |
| except Exception as exc: |
| return JSONResponse({"error": str(exc)}, status_code=500) |
|
|
| @app.post("/generate/stream") |
| async def generate_stream(req: GenerateRequest, request: Request): |
| auth_err = _check_api_key(request, settings) |
| if auth_err: |
| return auth_err |
|
|
| prompt, err = _validate_prompt(req.prompt, settings.max_prompt_length) |
| if err: |
| return err |
|
|
| if not model_service.model_exists: |
| return JSONResponse({"error": f"Model not found at {settings.model_path}"}, status_code=503) |
|
|
| sanitized = req.model_copy(update={"prompt": prompt}) |
| return StreamingResponse( |
| model_service.stream_tokens(sanitized), |
| media_type="text/event-stream", |
| headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}, |
| ) |
|
|
| @app.get("/info") |
| async def info(): |
| return JSONResponse({ |
| "name": APP_NAME, |
| "creator": CREATOR, |
| "version": APP_VERSION, |
| "status": "ready" if model_service.model_exists else "model-missing", |
| "mode": f"{'GPU' if settings.n_gpu_layers > 0 else 'CPU'} (llama.cpp)", |
| "gpu_layers": settings.n_gpu_layers, |
| "context_window": settings.n_ctx, |
| "model_loaded": model_service.model_loaded, |
| "load_time_seconds": model_service.load_time_seconds, |
| "endpoints": { |
| "generate": "POST /generate", |
| "stream": "POST /generate/stream", |
| "health": "GET /health", |
| "info": "GET /info", |
| }, |
| }) |
|
|
| return app |
|
|
|
|
| app = create_app() |
|
|
|
|
| def main() -> None: |
| s = app.state.settings |
| print(f"\n{'=' * 50}") |
| print(f"{APP_NAME.upper()} v{APP_VERSION}") |
| print(f"Creator: {CREATOR}") |
| print(f"{'=' * 50}") |
| print(f"Model: {s.model_path}") |
| print(f"GPU: {s.n_gpu_layers} layers") |
| print(f"Ctx: {s.n_ctx} | Threads: {s.n_threads} | Workers: {s.workers}") |
| print(f"URL: http://localhost:{s.port}\n") |
|
|
| uvicorn.run(app, host=s.host, port=s.port, log_level="warning") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
