helpme

#!/usr/bin/env python3
# AI shell helper
# By Dheera Venkatraman (https://dheera.net)

# Examples:
# $ helpme ls list all files in reverse order by size in human friendly units
# $ helpme dpkg remove packages with i386 architecture
# $ helpme apt install something to record my screen
# $ helpme docker run ubuntu 20.04, mount /home into it, and use the host networking
# $ helpme ffmpeg capture video from /dev/video0 every 1 second and write to .jpg files like img00000.jpg, img00001.jpg, ...
# $ helpme ffmpeg assemble all the .jpg files into an .mp4 timelapse video at 8fps
# $ helpme ffmpeg recompress myvideo.mp4 for HTML5-friendly use and save the result as myvideo_out.webm
# $ helpme convert Resize foo.jpg to 800x800 and center it on a 1200x1200 canvas with a white background and save as a .png file
# $ helpme gs convert somefile.pdf to a jpg files as page0.jpg, page1.jpg, etc. at 600dpi
# $ helpme inkscape export foo.svg to .eps and convert text to paths
# $ helpme ros2 publish a zero twist to /cmd_vel every 1 second
# $ helpme - rename all files in the current directory to lowercase

import sys
import os
import time

SERVICE = ""

# Ensure you have set your OpenAI or Anthropic API key as an environment variable
if os.getenv("OPENAI_API_KEY"):
    try:
        import openai
    except ImportError:
        # try harder, this is AI after all
        print("I found an OpenAI API key, but you need to run `pip3 install openai`.")
        if input("Should I run it for you (y/N)? ") == "y":
            os.system("pip3 install openai")
            import openai
        else:
            print("Aborted")
            exit(1)


    openai.api_key = os.getenv("OPENAI_API_KEY")
    client = openai.OpenAI()
    SERVICE = "openai"

elif os.getenv("ANTHROPIC_API_KEY"):
    try:
        import anthropic
    except ImportError:
        # try harder, this is AI after all
        print("I found an Anthropic API key, but you need to run `pip3 install anthropic`.")
        if input("Should I run it for you (y/N)? ") == "y":
            os.system("pip3 install anthropic")
            import anthropic
        else:
            print("Aborted")
            exit(1)
    client = anthropic.Client(api_key=os.getenv("ANTHROPIC_API_KEY"))
    SERVICE = "anthropic"

else:
    print("Error: Please set either OPENAI_API_KEY or ANTHROPIC_API_KEY. You can add it to .bashrc e.g.")
    print("export OPENAI_API_KEY=\"sk-....\"")
    print("export ANTHROPIC_API_KEY=\"sk-....\"")
    exit(1)

def get_command(executable, task_description):
    # Call the OpenAI API with the task description
    if executable == "-":
        system_role = f"You are an expert on the Linux command line and one-line shell scripts. Given the following English description of a task, you respond with the correct command. Do not include any other information in your response except for the command only. Do not enclose with triple backticks. If it cannot be done with common Linux command or utilities that you would usually find on a system, respond with the reason as a comment prefaced with #."
    else:
        system_role = f"You are an expert in `{executable}` commands on a Linux system. Given the following English description of a task, you respond with the correct `{executable}` command. Do not include any other information in your response except for the command only. Do not enclose with triple backticks. If `{executable}` is not known to you, `{executable}` cannot achieve the specified task, or you do not know the answer, respond with the reason as a comment prefaced with #."

    if SERVICE == "openai":
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role":"system", "content":system_role},
                {"role":"user", "content":task_description},
            ],
            max_tokens=150,
            temperature=0.0,
        )
        command = response.choices[0].message.content.strip()
    elif SERVICE == "anthropic":
        response = client.messages.create(
            model="claude-3-sonnet-20240229",
            system=system_role,
            messages=[
                {"role":"user", "content":task_description},
            ],
            max_tokens=150,
            temperature=0.0,
        )
        command = response.content[0].text
    return command

def main():
    if len(sys.argv) < 3:
        print("Usage: helpme <executable> <task_description>", file=sys.stderr)
        sys.exit(1)
    
    executable = sys.argv[1]
    task_description = " ".join(sys.argv[2:])
    command = get_command(executable, task_description)

    if command.startswith("#"):
        print(command.strip("# "), file=sys.stderr)
        exit(1)
    else:
        print(command)
        if input("Execute this command (y/N)? ") == "y":
            os.system(command)

if __name__ == "__main__":
    main()