llm.sh - a helper script wrapping llama.cpp

llm.sh

#!/usr/bin/env bash
#
# llm.sh - a thin wrapper around llama.cpp's llama-server for running
# reusable, named "tasks" (prompt templates) against a single local model.
#
# Usage:
#   llm server start|stop|status
#   llm tasks list
#   llm tasks new <name>
#   llm run <task_name> ["<text>"] [--auto-start] [--json]
#   cat file.txt | llm run <task_name> [--auto-start] [--json]
#
# Setup:
#   1. Save this file somewhere on your PATH, e.g. ~/bin/llm, and chmod +x it.
#   2. On first run it will scaffold ~/.llm/config.sh and ~/.llm/tasks/.
#   3. Edit ~/.llm/config.sh if you want a different model/port/etc.
#
# Requires: llama.cpp's `llama-server` binary on PATH, plus `curl` and `jq`.

set -euo pipefail

# ---------------------------------------------------------------------------
# Paths & defaults
# ---------------------------------------------------------------------------

LLM_HOME="${LLM_HOME:-$HOME/.llm}"
CONFIG_FILE="$LLM_HOME/config.sh"
TASKS_DIR="$LLM_HOME/tasks"

# Hardcoded fallback defaults (used to scaffold config.sh on first run, and
# as a safety net if a setting is missing from config.sh).
DEFAULT_HF_REPO="bartowski/Qwen2.5-3B-Instruct-GGUF:Q4_K_M"
DEFAULT_HOST="127.0.0.1"
DEFAULT_PORT="5555"
DEFAULT_CTX_SIZE="4096"
DEFAULT_PARALLEL="1"
DEFAULT_STARTUP_TIMEOUT="30"

# Per-task fallback sampling defaults (used when a task has no config.json).
DEFAULT_TEMPERATURE="30"
DEFAULT_MAX_TOKENS="50"
DEFAULT_TOP_P="1"

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

err() { echo "Error: $*" >&2; exit 1; }
info() { echo "$*" >&2; }

require_bin() {
  command -v "$1" >/dev/null 2>&1 || err "'$1' is required but not found on PATH."
}

ensure_home() {
  mkdir -p "$LLM_HOME" "$TASKS_DIR"

  if [[ ! -f "$CONFIG_FILE" ]]; then
    cat > "$CONFIG_FILE" <<EOF
# ~/.llm/config.sh - global settings for llm.sh
# Edit these values to change the model, port, context size, etc.

HF_REPO="$DEFAULT_HF_REPO"
HOST="$DEFAULT_HOST"
PORT="$DEFAULT_PORT"
CTX_SIZE="$DEFAULT_CTX_SIZE"
PARALLEL="$DEFAULT_PARALLEL"
STARTUP_TIMEOUT="$DEFAULT_STARTUP_TIMEOUT"

SERVER_LOG="\$LLM_HOME/server.log"
SERVER_PID_FILE="\$LLM_HOME/server.pid"
EOF
    info "Created default config at $CONFIG_FILE"
  fi
}

load_config() {
  ensure_home
  # shellcheck source=/dev/null
  source "$CONFIG_FILE"

  HF_REPO="${HF_REPO:-$DEFAULT_HF_REPO}"
  HOST="${HOST:-$DEFAULT_HOST}"
  PORT="${PORT:-$DEFAULT_PORT}"
  CTX_SIZE="${CTX_SIZE:-$DEFAULT_CTX_SIZE}"
  PARALLEL="${PARALLEL:-$DEFAULT_PARALLEL}"
  STARTUP_TIMEOUT="${STARTUP_TIMEOUT:-$DEFAULT_STARTUP_TIMEOUT}"
  SERVER_LOG="${SERVER_LOG:-$LLM_HOME/server.log}"
  SERVER_PID_FILE="${SERVER_PID_FILE:-$LLM_HOME/server.pid}"
}

base_url() { echo "http://${HOST}:${PORT}"; }

server_pid() {
  [[ -f "$SERVER_PID_FILE" ]] || return 1
  local pid
  pid="$(cat "$SERVER_PID_FILE" 2>/dev/null || true)"
  [[ -n "$pid" ]] || return 1
  kill -0 "$pid" 2>/dev/null || return 1
  echo "$pid"
}

server_healthy() {
  curl -s -o /dev/null -w '%{http_code}' "$(base_url)/health" 2>/dev/null | grep -q '^200$'
}

is_server_running() {
  server_pid >/dev/null 2>&1
}

# ---------------------------------------------------------------------------
# server subcommand
# ---------------------------------------------------------------------------

cmd_server_start() {
  if is_server_running; then
    info "Server already running (pid $(server_pid)) at $(base_url)"
    return 0
  fi

  require_bin llama-server

  info "Starting llama-server with model '$HF_REPO' on $(base_url) ..."
  nohup llama-server \
    -hf "$HF_REPO" \
    --host "$HOST" \
    --port "$PORT" \
    -c "$CTX_SIZE" \
    --parallel "$PARALLEL" \
    >> "$SERVER_LOG" 2>&1 &

  local pid=$!
  echo "$pid" > "$SERVER_PID_FILE"

  wait_for_server_ready || {
    err "Server did not become ready within ${STARTUP_TIMEOUT}s. Check $SERVER_LOG"
  }

  info "Server ready (pid $pid) at $(base_url). Logs: $SERVER_LOG"
}

wait_for_server_ready() {
  local waited=0
  while (( waited < STARTUP_TIMEOUT )); do
    if server_healthy; then
      return 0
    fi
    sleep 0.5
    waited=$(( waited + 1 ))  # loosely tracks half-seconds; timeout is generous
  done
  return 1
}

cmd_server_stop() {
  local pid
  if pid="$(server_pid)"; then
    info "Stopping server (pid $pid) ..."
    kill "$pid" 2>/dev/null || true
    rm -f "$SERVER_PID_FILE"
    info "Server stopped."
  else
    info "Server is not running."
  fi
}

cmd_server_status() {
  local pid
  if pid="$(server_pid)"; then
    if server_healthy; then
      echo "Running (pid $pid) at $(base_url) - healthy"
    else
      echo "Running (pid $pid) at $(base_url) - NOT responding to /health"
    fi
  else
    echo "Not running."
  fi
}

# ---------------------------------------------------------------------------
# tasks subcommand
# ---------------------------------------------------------------------------

cmd_tasks_list() {
  if [[ ! -d "$TASKS_DIR" ]] || [[ -z "$(ls -A "$TASKS_DIR" 2>/dev/null)" ]]; then
    info "No tasks yet. Create one with: llm tasks new <name>"
    return 0
  fi
  for dir in "$TASKS_DIR"/*/; do
    local name
    name="$(basename "$dir")"
    if [[ -f "$dir/system.txt" ]]; then
      echo "$name"
    fi
  done
}

cmd_tasks_new() {
  local name="${1:?Usage: llm tasks new <name>}"
  local dir="$TASKS_DIR/$name"

  [[ -e "$dir" ]] && err "Task '$name' already exists at $dir"

  mkdir -p "$dir"
  cat > "$dir/system.txt" <<'EOF'
You are a helpful assistant. Replace this with your task-specific instructions.
Respond with only the requested output, no extra commentary.
EOF

  cat > "$dir/config.json" <<EOF
{
  "temperature": $DEFAULT_TEMPERATURE,
  "max_tokens": $DEFAULT_MAX_TOKENS,
  "top_p": $DEFAULT_TOP_P
}
EOF

  info "Created task '$name' at $dir"
  info "Edit $dir/system.txt to define its prompt."
  info "Optionally add $dir/grammar.gbnf to constrain output."
}

# ---------------------------------------------------------------------------
# run subcommand
# ---------------------------------------------------------------------------

cmd_run() {
  local task_name="" input_text="" auto_start=false json_output=false
  local positional=()

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --auto-start) auto_start=true; shift ;;
      --json) json_output=true; shift ;;
      *) positional+=("$1"); shift ;;
    esac
  done

  [[ ${#positional[@]} -ge 1 ]] || err "Usage: llm run <task_name> [\"text\"] [--auto-start] [--json]"
  task_name="${positional[0]}"

  local task_dir="$TASKS_DIR/$task_name"
  [[ -f "$task_dir/system.txt" ]] || err "Unknown task '$task_name' (looked in $task_dir). Run: llm tasks list"

  if [[ ${#positional[@]} -ge 2 ]]; then
    input_text="${positional[1]}"
  else
    if [[ -t 0 ]]; then
      err "No input text given and no piped stdin. Usage: llm run $task_name \"text\", or pipe input via stdin."
    fi
    input_text="$(cat -)"
  fi

  if ! is_server_running; then
    if $auto_start; then
      cmd_server_start
    else
      err "Server not running. Run 'llm server start' first, or pass --auto-start."
    fi
  elif ! server_healthy; then
    err "Server process is running but not healthy. Check $SERVER_LOG"
  fi

  local system_prompt
  system_prompt="$(cat "$task_dir/system.txt")"

  # Merge task config.json (if present) over the global sampling defaults.
  local defaults_json task_json params_json
  defaults_json="$(jq -n \
    --argjson temperature "$DEFAULT_TEMPERATURE" \
    --argjson max_tokens "$DEFAULT_MAX_TOKENS" \
    --argjson top_p "$DEFAULT_TOP_P" \
    '{temperature: $temperature, max_tokens: $max_tokens, top_p: $top_p}')"

  if [[ -f "$task_dir/config.json" ]]; then
    task_json="$(cat "$task_dir/config.json")"
  else
    task_json='{}'
  fi

  params_json="$(jq -n --argjson a "$defaults_json" --argjson b "$task_json" '$a * $b')"

  # Build the request body, attaching a grammar file if the task has one.
  local body
  if [[ -f "$task_dir/grammar.gbnf" ]]; then
    local grammar
    grammar="$(cat "$task_dir/grammar.gbnf")"
    body="$(jq -n \
      --arg sys "$system_prompt" \
      --arg usr "$input_text" \
      --arg grammar "$grammar" \
      --argjson params "$params_json" \
      '{model: "local", messages: [{role: "system", content: $sys}, {role: "user", content: $usr}], grammar: $grammar} + $params')"
  else
    body="$(jq -n \
      --arg sys "$system_prompt" \
      --arg usr "$input_text" \
      --argjson params "$params_json" \
      '{model: "local", messages: [{role: "system", content: $sys}, {role: "user", content: $usr}]} + $params')"
  fi

  local response
  response="$(curl -s "$(base_url)/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d "$body")"

  if $json_output; then
    echo "$response"
  else
    echo "$response" | jq -r '.choices[0].message.content // ("Error: " + (.error.message // "unknown error"))'
  fi
}

# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

main() {
  require_bin curl
  require_bin jq
  load_config

  local cmd="${1:-}"
  [[ -n "$cmd" ]] || err "Usage: llm <server|tasks|run> ..."
  shift || true

  case "$cmd" in
    server)
      local sub="${1:-}"
      shift || true
      case "$sub" in
        start) cmd_server_start ;;
        stop) cmd_server_stop ;;
        status) cmd_server_status ;;
        *) err "Usage: llm server start|stop|status" ;;
      esac
      ;;
    tasks)
      local sub="${1:-}"
      shift || true
      case "$sub" in
        list) cmd_tasks_list ;;
        new) cmd_tasks_new "${1:-}" ;;
        *) err "Usage: llm tasks list|new <name>" ;;
      esac
      ;;
    run)
      cmd_run "$@"
      ;;
    *)
      err "Unknown command '$cmd'. Usage: llm <server|tasks|run> ..."
      ;;
  esac
}

main "$@"