nopilot.go

package main

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"

	openai "github.com/sashabaranov/go-openai"
)

const (
	BACKTICK       = "```"
	MODEL          = "gpt-3.5-turbo-0125"
	ASSISTANT_DESC = `Provide only code as output without any description.
Provide only code in plain text format without Markdown formatting.
Do not include symbols such as ` + BACKTICK + ` or ` + BACKTICK + `python.
If there is a lack of details, provide most logical solution.
You are not allowed to ask for more details.
For example if the prompt is "Hello world Python", you should return "print('Hello world')".`
)

var EDITOR = getEnv("EDITOR", "vim")

func getEnv(key, fallback string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return fallback
}

func promptEditor() string {
	tmpfile, err := ioutil.TempFile("", "*.txt")
	if err != nil {
		log.Fatalf("Error creating temporary file: %v", err)
	}
	defer os.Remove(tmpfile.Name())

	cmd := exec.Command(EDITOR, tmpfile.Name())
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		log.Fatalf("Error running editor command: %v", err)
	}

	content, err := os.ReadFile(tmpfile.Name())
	if err != nil {
		log.Fatalf("Error reading temporary file: %v", err)
	}

	return strings.TrimSpace(string(content))
}

func readInput() string {
	info, err := os.Stdin.Stat()
	if err != nil {
		log.Fatalf("Error stating stdin: %v", err)
	}

	if (info.Mode() & os.ModeCharDevice) == 0 {
		input, err := io.ReadAll(os.Stdin)
		if err != nil {
			log.Fatalf("Error reading stdin: %v", err)
		}
		return strings.TrimSpace(string(input))
	} else if len(os.Args) == 2 && os.Args[1] == "-e" {
		return promptEditor()
	} else if len(os.Args) > 1 {
		return strings.Join(os.Args[1:], " ")
	} else {
		fmt.Fprintln(os.Stderr, "No prompt given")
		os.Exit(1)
	}
	return ""
}

func initClient() *openai.Client {
	apiKey := os.Getenv("OPENAI_API_KEY")
	if apiKey == "" {
		log.Fatal("OPENAI_API_KEY environment variable not set")
	}
	return openai.NewClient(apiKey)
}

func genCode(client *openai.Client, prompt string) string {
	ctx := context.Background()
	req := openai.ChatCompletionRequest{
		Model: MODEL,
		Messages: []openai.ChatCompletionMessage{
			{Role: "system", Content: ASSISTANT_DESC},
			{Role: "user", Content: prompt},
		},
	}

	resp, err := client.CreateChatCompletion(ctx, req)
	if err != nil {
		log.Fatalf("Error generating code: %v", err)
	}

	return resp.Choices[0].Message.Content
}

func main() {
	client := initClient()
	prompt := readInput()
	text := genCode(client, prompt)
	fmt.Println("\n-------")
	fmt.Println(text)
}