learningcircuit--local-deep-research/cookiecutter-docker/hooks/pre_prompt.py

import json
import os
import subprocess
from typing import Any, Dict

import cookiecutter.prompt


def config_ollama(context: Dict[str, Any]) -> None:
    """
    Prompts the user for questions that are specific to Ollama. It is in a hook
    so that we can run it only if Ollama is enabled.

    """
    enable_ollama = cookiecutter.prompt.read_user_yes_no("enable_ollama", True)
    ollama_model = "gemma3:12b"
    if enable_ollama:
        # Ask ollama-specific questions.
        ollama_model = cookiecutter.prompt.read_user_variable(
            "ollama_model", ollama_model
        )

    context["_enable_ollama"] = enable_ollama
    context["_ollama_model"] = ollama_model


def check_gpu_linux(context: Dict[str, Any]) -> None:
    """
    Check if the system has an NVIDIA or AMD GPU on Linux.

    Args:
        context: The context dictionary to update with GPU information.

    """
    try:
        result = subprocess.run(["lspci"], capture_output=True, text=True)
        gpu_info = "\n".join(
            line for line in result.stdout.splitlines() if "vga" in line.lower()
        )
    except FileNotFoundError:
        gpu_info = ""

    if "NVIDIA" in gpu_info:
        print("Detected an Nvidia GPU.")
        context["_nvidia_gpu"] = True
        context["_amd_gpu"] = False
        context["enable_gpu"] = True
    elif "AMD" in gpu_info:
        print("Detected an AMD GPU.")
        context["_amd_gpu"] = True
        context["_nvidia_gpu"] = False
        context["enable_gpu"] = True
    else:
        print("Did not detect any GPU.")
        context["_nvidia_gpu"] = False
        context["_amd_gpu"] = False
        context["enable_gpu"] = False


def check_gpu_windows(context: Dict[str, Any]) -> None:
    """
    Check if the system has an NVIDIA or AMD GPU on Windows.

    Args:
        context: The context dictionary to update with GPU information.

    """
    result = subprocess.run(
        ["wmic", "path", "win32_VideoController", "get", "name"],
        capture_output=True,
        text=True,
    )
    gpu_info = result.stdout.strip()

    if "NVIDIA" in gpu_info.upper():
        print("Detected an Nvidia GPU.")
        context["_nvidia_gpu"] = True
        context["_amd_gpu"] = False
        context["enable_gpu"] = True
    elif "AMD" in gpu_info.upper() or "RADEON" in gpu_info.upper():
        print("Detected an AMD GPU.")
        context["_amd_gpu"] = True
        context["_nvidia_gpu"] = False
        context["enable_gpu"] = True
    else:
        print("Did not detect any GPU.")
        context["_nvidia_gpu"] = False
        context["_amd_gpu"] = False
        context["enable_gpu"] = False


def main() -> None:
    # Load the context.
    with open("cookiecutter.json", "r", encoding="utf-8") as config:
        context = json.load(config)

    # Check GPU information based on the operating system.
    if os.name == "posix" and os.uname().sysname == "Linux":
        check_gpu_linux(context)
    elif os.name == "nt":
        check_gpu_windows(context)
    # Ollama-specific config.
    config_ollama(context)

    # Save the updated context back to cookiecutter.json.
    with open("cookiecutter.json", "w", encoding="utf-8") as config:
        json.dump(context, config, indent=4)


if __name__ == "__main__":
    main()