zbot.py

#!/usr/bin/env python3
"""ZBOT - ROZ AI Assistant with auto hardware detection and model selection."""

import subprocess
import json
import os
import sys
import re

MASCOT = r"""
    \033[1;31m╔══════════════════════╗
    ║   ┌─────────────┐   ║
    ║   │  ◉       ◉  │   ║
    ║   │     ───     │   ║
    ║   │   \_____/   │   ║
    ║   └─────────────┘   ║
    ║      ┌───────┐      ║
    ║      │ Z-BOT │      ║
    ║      └───────┘      ║
    ╚══════════════════════╝\033[0m"""

MODEL_TIERS = [
    {"min_vram": 16000, "min_ram": 32, "model": "llama3.1:70b-instruct-q4_0", "label": "Llama 3.1 70B (beast mode)"},
    {"min_vram": 8000,  "min_ram": 16, "model": "llama3.1:8b",               "label": "Llama 3.1 8B"},
    {"min_vram": 4000,  "min_ram": 8,  "model": "llama3.2",                  "label": "Llama 3.2 3B"},
    {"min_vram": 2000,  "min_ram": 4,  "model": "llama3.2:1b",              "label": "Llama 3.2 1B (lightweight)"},
    {"min_vram": 0,     "min_ram": 2,  "model": "tinyllama",                "label": "TinyLlama (minimal)"},
]

SYSTEM_PROMPT = """You are ZBOT, the AI assistant for ROZ OS (Republic of Zani).
You are helpful, direct, and a bit cheeky. You can help with anything on the user's Linux system.
When the user asks you to run a command, output it in this format: [RUN] command here
Keep responses short and useful. You're running locally on their hardware — no cloud, fully private."""


def run(cmd):
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=10)
        return r.stdout.strip()
    except Exception:
        return ""


def detect_hardware():
    hw = {"gpu": None, "vram_mb": 0, "ram_gb": 0, "cpu": "Unknown"}
    # GPU
    gpu = run("nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null")
    if gpu:
        parts = gpu.split(",")
        hw["gpu"] = parts[0].strip()
        vram = re.search(r'(\d+)', parts[1]) if len(parts) > 1 else None
        hw["vram_mb"] = int(vram.group(1)) if vram else 0
    else:
        amd = run("lspci | grep -i 'VGA.*AMD' 2>/dev/null")
        if amd:
            hw["gpu"] = "AMD GPU"
            hw["vram_mb"] = 4000  # estimate
    # RAM
    ram = run("free -g | awk '/Mem/{print $2}'")
    hw["ram_gb"] = int(ram) if ram.isdigit() else 4
    # CPU
    hw["cpu"] = run("lscpu | grep 'Model name' | sed 's/.*: *//'") or "Unknown"
    return hw


def pick_model(hw):
    for tier in MODEL_TIERS:
        if hw["vram_mb"] >= tier["min_vram"] and hw["ram_gb"] >= tier["min_ram"]:
            return tier["model"], tier["label"]
    return "tinyllama", "TinyLlama (fallback)"


def ensure_model(model):
    installed = run("ollama list 2>/dev/null")
    if model.split(":")[0] in installed:
        return True
    print(f"\033[1;33m⬇️  Downloading {model}...\033[0m")
    os.system(f"ollama pull {model}")
    return True


def execute_command(cmd):
    print(f"\033[1;33m⚡ Running: {cmd}\033[0m")
    os.system(cmd)


def chat(model):
    history = []
    while True:
        try:
            user_input = input("\n\033[1;37mYou > \033[0m").strip()
        except (KeyboardInterrupt, EOFError):
            print("\n\033[1;31m👋 ZBOT out. Stay sharp.\033[0m")
            break
        if not user_input:
            continue
        if user_input.lower() in ("exit", "quit", "bye"):
            print("\033[1;31m👋 ZBOT out. Stay sharp.\033[0m")
            break

        history.append({"role": "user", "content": user_input})
        try:
            result = subprocess.run(
                ["ollama", "run", model, "--nowordwrap"],
                input=f"System: {SYSTEM_PROMPT}\n\n" + "\n".join(
                    f"{'User' if m['role']=='user' else 'ZBOT'}: {m['content']}" for m in history[-10:]
                ),
                capture_output=True, text=True, timeout=120
            )
            response = result.stdout.strip()
        except subprocess.TimeoutExpired:
            response = "Sorry, that took too long. Try again?"
        except Exception as e:
            response = f"Error: {e}"

        # Check for commands to execute
        if "[RUN]" in response:
            lines = response.split("\n")
            for line in lines:
                if "[RUN]" in line:
                    cmd = line.split("[RUN]")[1].strip()
                    clean_response = response.replace(line, "").strip()
                    if clean_response:
                        print(f"\033[1;31mZBOT > \033[0m{clean_response}")
                    confirm = input(f"\033[1;33m⚡ Run '{cmd}'? [y/n] \033[0m").strip().lower()
                    if confirm == "y":
                        execute_command(cmd)
                else:
                    if line.strip():
                        print(f"\033[1;31mZBOT > \033[0m{line}")
        else:
            print(f"\033[1;31mZBOT > \033[0m{response}")

        history.append({"role": "assistant", "content": response})


def main():
    print(MASCOT)
    print("\033[1;31m  ZBOT v1.0 — ROZ AI Assistant\033[0m")
    print("\033[0;37m  100% local • 100% private • 0% cloud\033[0m\n")

    hw = detect_hardware()
    print(f"  🖥️  CPU:  {hw['cpu']}")
    print(f"  🎮 GPU:  {hw['gpu'] or 'None'} ({hw['vram_mb']}MB VRAM)")
    print(f"  🧠 RAM:  {hw['ram_gb']}GB")

    model, label = pick_model(hw)
    print(f"  🤖 Model: {label}")
    print(f"  📦 Using: {model}\n")

    ensure_model(model)
    print("\033[1;32m  ✅ Ready! Type 'exit' to quit.\033[0m")
    chat(model)


if __name__ == "__main__":
    main()