augment17/claude-code-backend

1# -*- coding: utf-8 -*-2"""3swarm_llm.py — Local CPU LLM Swarm (Qwen2.5-1.5B-Instruct, Q4_K_M)4━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━5Why a local model?6  The NIM API is rate-limited (tokens/minute). Every small sub-task7  (JSON formatting, log summarization, brainstorm question generation,8  Bell Curve trimming) that goes to NIM wastes quota needed for coding.9 10This module runs Qwen2.5-1.5B-Instruct at Q4_K_M quantization:11  - RAM:   ~1.1 GB (leaves 14 GB free for FastAPI + data)12  - Speed: ~45 tok/s on 2 vCPUs (good enough for short tasks)13  - Model: downloaded from HF Hub on first run → cached in /tmp/models/14 15NIM is used ONLY for heavy coding tasks (forge_execute).16SwarmLLM handles everything else.17 18If llama-cpp-python is not installed (e.g. first boot before pip):19  the module silently degrades and returns a FALLBACK_STUB response,20  so the rest of the system still works.21"""22 23import os24import logging25import asyncio26import time27from typing import Optional28 29logger = logging.getLogger("swarm_llm")30 31MODEL_CACHE_DIR  = os.environ.get("MODEL_CACHE_DIR", "/tmp/models")32MODEL_REPO       = "Qwen/Qwen2.5-1.5B-Instruct-GGUF"33MODEL_FILENAME   = "qwen2.5-1.5b-instruct-q4_k_m.gguf"34N_CTX            = 2048   # context window — matches our Bell Curve budget35N_THREADS        = int(os.environ.get("SWARM_THREADS", "2"))36MAX_TOKENS       = int(os.environ.get("SWARM_MAX_TOKENS", "256"))37ENABLE_SWARM     = os.environ.get("ENABLE_SWARM_LLM", "true").lower() == "true"38 39_llm = None          # loaded lazily on first call40_llm_lock = None     # asyncio.Lock initialised at first call41 42 43def _get_lock():44    global _llm_lock45    if _llm_lock is None:46        _llm_lock = asyncio.Lock()47    return _llm_lock48 49 50def _load_model() -> Optional[object]:51    """Download (if needed) and load the GGUF model. Blocking — call in executor."""52    global _llm53    if _llm is not None:54        return _llm55    if not ENABLE_SWARM:56        logger.info("[SwarmLLM] Disabled via ENABLE_SWARM_LLM=false")57        return None58    try:59        from llama_cpp import Llama60    except ImportError:61        logger.warning("[SwarmLLM] llama-cpp-python not installed. Running in stub mode.")62        return None63 64    model_path = os.path.join(MODEL_CACHE_DIR, MODEL_FILENAME)65    if not os.path.exists(model_path):66        logger.info(f"[SwarmLLM] Downloading {MODEL_FILENAME} from HF Hub…")67        try:68            from huggingface_hub import hf_hub_download69            os.makedirs(MODEL_CACHE_DIR, exist_ok=True)70            model_path = hf_hub_download(71                repo_id=MODEL_REPO,72                filename=MODEL_FILENAME,73                local_dir=MODEL_CACHE_DIR,74                local_dir_use_symlinks=False,75            )76            logger.info(f"[SwarmLLM] Downloaded → {model_path}")77        except Exception as e:78            logger.error(f"[SwarmLLM] Download failed: {e}")79            return None80 81    logger.info(f"[SwarmLLM] Loading model (n_ctx={N_CTX}, threads={N_THREADS}, type_k=8 (Q8_0), type_v=8 (Q8_0))…")82    t0 = time.time()83    try:84        _llm = Llama(85            model_path=model_path,86            n_ctx=N_CTX,87            n_threads=N_THREADS,88            n_gpu_layers=0,      # CPU only — HF free spaces have no GPU89            verbose=False,90            chat_format="chatml",91            type_k=8,            # 8-bit quantization for Key Cache (Turbo Quant)92            type_v=8,            # 8-bit quantization for Value Cache (Turbo Quant)93        )94        logger.info(f"[SwarmLLM] Model loaded in {time.time()-t0:.1f}s")95        return _llm96    except Exception as e:97        logger.error(f"[SwarmLLM] Failed to load model: {e}")98        return None99 100 101class SwarmLLM:102    """103    Async wrapper around the local Qwen model.104    All inference runs in a thread executor so the FastAPI event loop105    is never blocked.106    """107 108    def __init__(self):109        self._ready = False110 111    async def warm_up(self):112        """Pre-load the model at startup so first inference is instant."""113        loop = asyncio.get_event_loop()114        model = await loop.run_in_executor(None, _load_model)115        self._ready = model is not None116        if self._ready:117            logger.info("[SwarmLLM] Warm-up complete. Ready for inference.")118        else:119            logger.warning("[SwarmLLM] Running in stub mode (model not available).")120 121    async def infer(self, prompt: str, system: str = "", max_tokens: int = MAX_TOKENS) -> str:122        """123        Run inference on the local model. Returns the generated text.124        Falls back to a stub if model is not loaded.125        """126        if not self._ready:127            return self._stub(prompt)128 129        lock = _get_lock()130        async with lock:131            loop = asyncio.get_event_loop()132            result = await loop.run_in_executor(133                None,134                lambda: self._sync_infer(prompt, system, max_tokens),135            )136        return result137 138    def _sync_infer(self, prompt: str, system: str, max_tokens: int) -> str:139        global _llm140        if _llm is None:141            return self._stub(prompt)142        try:143            messages = []144            if system:145                messages.append({"role": "system", "content": system})146            messages.append({"role": "user", "content": prompt})147 148            response = _llm.create_chat_completion(149                messages=messages,150                max_tokens=max_tokens,151                temperature=0.3,152                stop=["<|im_end|>", "</s>"],153            )154            text = response["choices"][0]["message"]["content"].strip()155            logger.debug("[SwarmLLM] Infer complete: %d chars", len(text))156            return text157        except Exception as e:158            logger.error(f"[SwarmLLM] Inference error: {e}")159            return self._stub(prompt)160 161    @staticmethod162    def _stub(prompt: str) -> str:163        """Fallback when model is unavailable — returns a safe placeholder."""164        return "[SwarmLLM unavailable — NIM will handle this task]"165 166    # ── High-Level Task Shortcuts ──────────────────────────────────────────────167 168    async def summarize(self, text: str, max_words: int = 80) -> str:169        """Summarise a long text into ≤ max_words words. Used to enforce Bell Curve budget."""170        prompt = (171            f"Summarise the following in ≤ {max_words} words. "172            f"Be dense with information. No filler sentences.\n\n{text[:3000]}"173        )174        return await self.infer(prompt, system="You are a precise technical summariser.")175 176    async def format_json(self, raw: str) -> str:177        """Extract and clean a JSON object from a messy LLM response."""178        prompt = (179            "Extract the JSON object from the following text. "180            "Return ONLY the JSON, no markdown fences, no explanation.\n\n" + raw[:2000]181        )182        return await self.infer(prompt, system="You are a JSON extractor. Output only valid JSON.")183 184    async def generate_brainstorm_questions(self, goal: str, n: int = 5) -> list:185        """Generate n 'What if…' brainstorm questions for the hourly swarm cycle."""186        prompt = (187            f"Generate exactly {n} creative 'What if…' questions to improve: '{goal}'. "188            f"Each question should be a novel technical idea. "189            f"Format: one question per line, no numbering."190        )191        raw = await self.infer(prompt, system="You are a creative technical brainstormer.", max_tokens=400)192        questions = [q.strip() for q in raw.strip().splitlines() if q.strip()]193        return questions[:n]194 195    async def smart_trim(self, text: str, char_limit: int) -> str:196        """197        If text exceeds char_limit, ask the local model to summarise it198        to fit. Better than hard-cutting. Used in Bell Curve budget enforcement.199        """200        if len(text) <= char_limit:201            return text202        target_words = char_limit // 6  # rough chars-to-words ratio203        summary = await self.summarize(text, max_words=target_words)204        return summary205 206 207# Singleton — import and use directly208swarm = SwarmLLM()209