llama-cpp-wrapper.py

# /// script
# requires-python = ">=3.8"
# dependencies = [
#     "llama-cpp-python",
# ]
# ///

"""
A basic llama.cpp wrapper to explore LLMs in an interactive chat session with channel support and reasoning.
"""

import argparse
import os
import sys
import time
import threading
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from llama_cpp import Llama
try:
    from llama_cpp import llama_log_set  # may not exist on older wheels
except Exception:
    llama_log_set = None

from contextlib import contextmanager

@contextmanager
def _suppress_c_stderr():
    """Temporarily silence C-level stderr (fd=2). Used around model load."""
    try:
        fd = sys.stderr.fileno()
    except Exception:
        yield
        return
    saved = os.dup(fd)
    devnull = os.open(os.devnull, os.O_WRONLY)
    try:
        os.dup2(devnull, fd)
        os.close(devnull)
        yield
    finally:
        try:
            os.dup2(saved, fd)
        finally:
            os.close(saved)

def _install_quiet_llama_logger(enable: bool) -> None:
    if not enable or llama_log_set is None:
        return
    try:
        def _drop_logs(level, text):
            return None
        llama_log_set(_drop_logs)
    except Exception:
        pass

TAG_OPEN = "<|"
TAG_CLOSE = "|>"

REASONING_CHANNELS = {"analysis", "critic", "plan", "thought", "trace"}
FINAL_CHANNELS = {"final", "answer"}

CHANNEL_COLORS: Dict[str, str] = {
    "analysis": "90",  # dim gray
    "critic": "33",    # yellow
    "plan": "36",      # cyan
    "thought": "35",   # magenta
    "trace": "94",     # bright blue
    "final": "97",     # bright white
    "answer": "97",    # bright white
    "tool": "92",      # bright green
}
DEFAULT_COLOR = "37"

SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]


def color_for_channel(ch: Optional[str]) -> str:
    if not ch:
        return DEFAULT_COLOR
    return CHANNEL_COLORS.get(ch.lower(), DEFAULT_COLOR)


def print_colored(text: str, color_code: str) -> None:
    if not text:
        return
    sys.stdout.write(f"\033[{color_code}m{text}\033[0m")
    sys.stdout.flush()


@dataclass
class ParseEvent:
    kind: str  # "role" | "channel" | "text"
    value: str
    role: Optional[str] = None
    channel: Optional[str] = None


class LlamaChannelParser:
    def __init__(self) -> None:
        self.buf: str = ""
        self.role: Optional[str] = None
        self.channel: Optional[str] = None
        self.mode: str = "outside"  # "outside" | "message" | "read_name"
        self.pending_kind: Optional[str] = None
        self.name_buf: str = ""

    def feed(self, chunk: str) -> List[ParseEvent]:
        self.buf += chunk
        out: List[ParseEvent] = []
        while True:
            if self.mode == "message":
                idx = self.buf.find(TAG_OPEN)
                if idx == -1:
                    if self.buf:
                        out.append(ParseEvent("text", self.buf, self.role, self.channel))
                        self.buf = ""
                    break
                if idx > 0:
                    out.append(ParseEvent("text", self.buf[:idx], self.role, self.channel))
                    self.buf = self.buf[idx:]
                self.mode = "outside"
                continue

            if self.mode == "read_name":
                idx = self.buf.find(TAG_OPEN)
                if idx == -1:
                    self.name_buf += self.buf
                    self.buf = ""
                    break
                self.name_buf += self.buf[:idx]
                self.buf = self.buf[idx:]
                name = self.name_buf.strip()
                self.name_buf = ""
                kind = self.pending_kind
                self.pending_kind = None
                self.mode = "outside"
                if kind == "start":
                    self.role = name or None
                    out.append(ParseEvent("role", self.role or ""))
                elif kind == "channel":
                    self.channel = name or None
                    out.append(ParseEvent("channel", self.channel or ""))
                continue

            if not self.buf:
                break
            if not self.buf.startswith(TAG_OPEN):
                idx = self.buf.find(TAG_OPEN)
                if idx == -1:
                    break
                self.buf = self.buf[idx:]

            close = self.buf.find(TAG_CLOSE, len(TAG_OPEN))
            if close == -1:
                break
            token = self.buf[len(TAG_OPEN):close].strip()
            self.buf = self.buf[close + len(TAG_CLOSE):]

            if token == "message":
                self.mode = "message"
                continue
            elif token in ("start", "channel"):
                self.pending_kind = token
                self.mode = "read_name"
                continue
            else:
                continue
        return out


def parse_args():
    p = argparse.ArgumentParser(description="GGUF Model Chat Interface")
    p.add_argument("-d", "--debug", action="store_true", help="Enable debug output")
    p.add_argument("-m", "--model", required=True, help="Path to GGML model file, typically .bin or .gguf")
    p.add_argument("-c", "--context", type=int, default=32768, help="Context window size")
    p.add_argument("-t", "--threads", type=int, help="Number of threads (default: half of CPU cores)")

    g = p.add_mutually_exclusive_group()
    g.add_argument("--simple-output", dest="simple_output", action="store_true", help="Enable simple output (spinners for non-main channels)")
    g.add_argument("--no-simple-output", dest="simple_output", action="store_false", help="Disable simple output; print all channels verbatim with prefixes for non-main channels")
    p.set_defaults(simple_output=True)

    p.add_argument("--system-prompt", help="System prompt text to prepend as the first message")
    p.add_argument("--system-prompt-file", help="Path to a file containing the system prompt text")

    # Suppress llama.cpp / ggml logs during model load (default on unless --no-quiet-llama)
    p.add_argument("--no-quiet-llama", dest="quiet_llama", action="store_false",
                   help="Do not suppress llama.cpp/ggml stderr and logs during model load")
    p.set_defaults(quiet_llama=True)

    return p.parse_args()


def load_model(model_path: str, context_size: int, threads: int, debug: bool, quiet: bool) -> Llama:
    if not os.path.exists(model_path):
        print(f"Error: Model file '{model_path}' not found")
        sys.exit(1)
    if debug:
        print(f"Loading model: {model_path}")
        print(f"Context size: {context_size}")
        print(f"Threads: {threads}")

    _install_quiet_llama_logger(enable=(not debug and quiet))

    try:
        if not debug and quiet:
            with _suppress_c_stderr():
                return Llama(model_path=model_path, n_threads=threads, n_ctx=context_size, verbose=debug)
        else:
            return Llama(model_path=model_path, n_threads=threads, n_ctx=context_size, verbose=debug)
    except Exception as e:
        print(f"Error loading model: {e}")
        sys.exit(1)


def generate_simple_fallback(llm: Llama, messages: List[Dict[str, str]], temperature=0.7, max_tokens=1024, debug=False) -> str:
    prompt_parts = []
    for msg in messages:
        role = msg.get("role", "")
        content = msg.get("content", "")
        if role == "user":
            prompt_parts.append(f"User: {content}")
        elif role == "assistant":
            prompt_parts.append(f"Assistant: {content}")
    prompt = "\n".join(prompt_parts) + "\nAssistant:"

    try:
        resp = llm(
            prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=0.9,
            repeat_penalty=1.1,
            stop=["User:", "\nUser:"],
            stream=True,
        )
        out = ""
        for o in resp:
            if 'choices' in o and o['choices']:
                tok = o['choices'][0].get('text')
                if tok:
                    print_colored(tok, DEFAULT_COLOR)
                    out += tok
        return out.strip()
    except Exception as e:
        print(f"\nGeneration error: {e}")
        return ""


def generate_response(
    llm: Llama,
    messages: List[Dict[str, str]],
    simple_output: bool = True,
    temperature: float = 0.7,
    max_tokens: int = 1024,
    debug: bool = False,
) -> str:
    start_time = time.time()
    first_token_time = None

    if debug:
        print(f"\n[DEBUG] Messages: {len(messages)} in history")
        print(f"[DEBUG] Simple output: {simple_output}")
        print(f"[DEBUG] Generation params: temp={temperature}, max_tokens={max_tokens}")

    # Waiting spinner (unlabeled, no preceding newline)
    wait_stop = threading.Event()

    def wait_spinner():
        idx = 0
        last_len = 0
        while not wait_stop.is_set():
            frame = SPINNER_FRAMES[idx]
            sys.stdout.write("\r" + frame)
            if last_len > len(frame):
                sys.stdout.write(" " * (last_len - len(frame)))
                sys.stdout.write("\r" + frame)
            sys.stdout.flush()
            last_len = len(frame)
            idx = (idx + 1) % len(SPINNER_FRAMES)
            time.sleep(0.1)
        sys.stdout.write("\r" + (" " * last_len) + "\r")
        sys.stdout.flush()

    waiter = threading.Thread(target=wait_spinner, daemon=True)
    waiter.start()

    try:
        response = llm.create_chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=0.9,
            repeat_penalty=1.1,
            stream=True,
        )

        parser = LlamaChannelParser()
        collected: Dict[Tuple[Optional[str], Optional[str]], str] = {}
        printed: Dict[Tuple[Optional[str], Optional[str]], int] = {}
        prefix_printed: Dict[Tuple[Optional[str], Optional[str]], bool] = {}
        token_count = 0

        last_was_newline = True

        active_spinner_key: Optional[Tuple[Optional[str], Optional[str]]] = None
        active_spinner_channel: Optional[str] = None
        chan_spinner_idx = 0

        def spinner_start(channel_name: str):
            nonlocal chan_spinner_idx, last_was_newline
            chan_spinner_idx = 0
            col = color_for_channel(channel_name)
            sys.stdout.write("\r")
            print_colored(f"[{channel_name}] {SPINNER_FRAMES[chan_spinner_idx]}", col)
            sys.stdout.write("\r")
            sys.stdout.flush()
            last_was_newline = False

        def spinner_tick(channel_name: str):
            nonlocal chan_spinner_idx, last_was_newline
            chan_spinner_idx = (chan_spinner_idx + 1) % len(SPINNER_FRAMES)
            col = color_for_channel(channel_name)
            sys.stdout.write("\r")
            print_colored(f"[{channel_name}] {SPINNER_FRAMES[chan_spinner_idx]}", col)
            sys.stdout.flush()
            last_was_newline = False

        def spinner_done(channel_name: str):
            nonlocal last_was_newline
            col = color_for_channel(channel_name)
            sys.stdout.write("\r")
            print_colored(f"[{channel_name}] done", col)
            sys.stdout.write("\n")
            sys.stdout.flush()
            last_was_newline = True

        def close_active_spinner_if_any():
            nonlocal active_spinner_key, active_spinner_channel
            if active_spinner_key is not None and active_spinner_channel is not None:
                spinner_done(active_spinner_channel)
                active_spinner_key = None
                active_spinner_channel = None

        for chunk in response:
            if 'choices' not in chunk or not chunk['choices']:
                continue
            delta = chunk['choices'][0].get('delta', {})
            piece = delta.get('content', "")
            if not piece:
                continue

            if first_token_time is None:
                first_token_time = time.time()
                wait_stop.set()
                waiter.join(timeout=1.0)
            token_count += 1

            events = parser.feed(piece)
            for ev in events:
                if ev.kind == "text":
                    key = (ev.role, ev.channel)
                    prev = collected.get(key, "")
                    collected[key] = prev + ev.value

                    ch = (ev.channel or "").lower()
                    col = color_for_channel(ch)
                    is_main = (ch in FINAL_CHANNELS) or (ch == "")

                    if simple_output and not is_main:
                        if active_spinner_key != key:
                            close_active_spinner_if_any()
                            active_spinner_key = key
                            active_spinner_channel = ch or "other"
                            spinner_start(active_spinner_channel)
                        else:
                            spinner_tick(active_spinner_channel or "other")
                    else:
                        if is_main:
                            close_active_spinner_if_any()
                        already = printed.get(key, 0)
                        new_text = collected[key][already:]
                        if new_text:
                            if (not simple_output) and (not is_main) and not prefix_printed.get(key, False):
                                if not last_was_newline:
                                    sys.stdout.write("\n")
                                print_colored(f"[{ch}] ", col)
                                prefix_printed[key] = True
                                last_was_newline = False
                            if already == 0 and is_main and not last_was_newline:
                                sys.stdout.write("\n")
                            print_colored(new_text, col)
                            printed[key] = already + len(new_text)
                            last_was_newline = new_text.endswith("\n")

                elif ev.kind in ("role", "channel"):
                    close_active_spinner_if_any()

        close_active_spinner_if_any()
        if first_token_time is None:
            wait_stop.set()
            waiter.join(timeout=1.0)

        def pick_final() -> str:
            order = [("assistant", "final"), ("assistant", "answer"), ("assistant", None)]
            for rk, ck in order:
                for (r, c), txt in collected.items():
                    if r == rk and ((ck is None and (c is None or c == "")) or (c == ck)):
                        return txt
            parts = []
            for (r, c), txt in collected.items():
                if r == "assistant":
                    ch = (c or "").lower()
                    if ch not in REASONING_CHANNELS:
                        parts.append(txt)
            return "".join(parts).strip()

        final_text = pick_final()

        if debug and first_token_time:
            elapsed = time.time() - start_time
            first_token_latency = first_token_time - start_time
            tps = token_count / elapsed if elapsed > 0 else 0
            print(f"\n[DEBUG] First token: {first_token_latency:.2f}s")
            print(f"[DEBUG] Total time: {elapsed:.2f}s")
            print(f"[DEBUG] Tokens: {token_count}")
            by_rc = {f"{r or ''}:{c or ''}": len(txt) for (r, c), txt in collected.items()}
            print(f"[DEBUG] By (role:channel): {by_rc}")

        return final_text.strip()

    except KeyboardInterrupt:
        wait_stop.set()
        try:
            waiter.join(timeout=1.0)
        except Exception:
            pass
        print("\n[exit]")
        return ""
    except Exception as e:
        wait_stop.set()
        try:
            waiter.join(timeout=1.0)
        except Exception:
            pass
        return generate_simple_fallback(llm, messages, temperature, max_tokens, debug)


def main() -> None:
    args = parse_args()
    threads = args.threads or (os.cpu_count() // 2)
    llm = load_model(args.model, args.context, threads, args.debug, args.quiet_llama)

    print("Interactive llama.cpp Chat Session")
    print("Commands: /exit, /clear, /help, /simple")
    if args.debug:
        print("[Debug mode enabled]")
    print(f"[Simple output {'enabled' if args.simple_output else 'disabled'}]")
    print()

    messages: List[Dict[str, str]] = []

    if args.system_prompt and args.system_prompt_file:
        print("Error: specify either --system-prompt or --system-prompt-file, not both.")
        sys.exit(1)

    system_prompt_text: Optional[str] = None
    if args.system_prompt:
        system_prompt_text = args.system_prompt
    elif args.system_prompt_file:
        if not os.path.exists(args.system_prompt_file):
            print(f"Error: system prompt file '{args.system_prompt_file}' not found")
            sys.exit(1)
        with open(args.system_prompt_file, "r", encoding="utf-8") as f:
            system_prompt_text = f.read()

    if system_prompt_text:
        messages.append({"role": "system", "content": system_prompt_text})
        if args.debug:
            print("[DEBUG] Loaded system prompt (", len(system_prompt_text), "chars )")

    simple_output = args.simple_output

    while True:
        try:
            user_input = input("❯ ").strip()
            if not user_input:
                continue

            low = user_input.lower()
            if low == "/exit":
                print("[exit]")
                break
            elif low == "/clear":
                messages = []
                if system_prompt_text:
                    messages.append({"role": "system", "content": system_prompt_text})
                print("Chat history cleared.")
                continue
            elif low == "/simple":
                simple_output = not simple_output
                print(f"Simple output {'enabled' if simple_output else 'disabled'}.")
                continue
            elif low == "/help":
                print("Commands:")
                print("  /exit    - Exit the chat")
                print("  /clear   - Clear chat history (re-adds system prompt if set)")
                print("  /simple  - Toggle simple output (spinners for non-main channels)")
                print("  /help    - Show this help")
                continue

            messages.append({"role": "user", "content": user_input})
            response = generate_response(llm, messages, simple_output=simple_output, debug=args.debug)

            if response:
                print()
                messages.append({"role": "assistant", "content": response})

        except KeyboardInterrupt:
            print("\n[exit]")
            break
        except EOFError:
            print("\n[exit]")
            break


if __name__ == "__main__":
    main()