x

Programs, configuration and documentation that don't fit anywhere else
Log | Files | Refs | README | LICENSE

llm.go (2419B)

      1 package main
      2 
      3 import (
      4 	"bufio"
      5 	"bytes"
      6 	"errors"
      7 	"flag"
      8 	"fmt"
      9 	"io"
     10 	"log"
     11 	"net/http"
     12 	"os"
     13 	"path"
     14 
     15 	"olowe.co/x/internal/aiconf"
     16 	"olowe.co/x/openai"
     17 )
     18 
     19 var model = flag.String("m", "", "model")
     20 var sysPrompt = flag.String("s", "", "system prompt")
     21 var converse = flag.Bool("c", false, "start a back-and-forth chat")
     22 
     23 func copyAll(w io.Writer, paths []string) (n int64, err error) {
     24 	if len(paths) == 0 {
     25 		return io.Copy(w, os.Stdin)
     26 	}
     27 	var errs []error
     28 	for _, name := range paths {
     29 		f, err := os.Open(name)
     30 		if err != nil {
     31 			return n, err
     32 		}
     33 		nn, err := io.Copy(w, f)
     34 		if err != nil {
     35 			errs = append(errs, fmt.Errorf("copy %s: %w", name, err))
     36 		}
     37 		f.Close()
     38 		n += nn
     39 	}
     40 	return n, errors.Join(errs...)
     41 }
     42 
     43 func init() {
     44 	log.SetFlags(0)
     45 	log.SetPrefix("llm: ")
     46 	flag.Parse()
     47 }
     48 
     49 func main() {
     50 	confDir, err := os.UserConfigDir()
     51 	if err != nil {
     52 		log.Fatal(err)
     53 	}
     54 	config, err := aiconf.ReadConfig(path.Join(confDir, "openai"))
     55 	if err != nil {
     56 		log.Fatalf("read configuration: %v", err)
     57 	}
     58 	if *model == "" {
     59 		*model = config.DefaultModel
     60 	}
     61 	if config.BaseURL == "" {
     62 		config.BaseURL = "http://127.0.0.1:8080"
     63 	}
     64 	client := &openai.Client{http.DefaultClient, config.Token, config.BaseURL}
     65 
     66 	chat := openai.Chat{Model: *model}
     67 	if *sysPrompt != "" {
     68 		chat.Messages = []openai.Message{
     69 			{
     70 				Role:    openai.RoleSystem,
     71 				Content: *sysPrompt,
     72 			},
     73 		}
     74 	}
     75 	buf := &bytes.Buffer{}
     76 	if !*converse {
     77 		_, err := copyAll(buf, flag.Args())
     78 		if err != nil {
     79 			log.Fatalln("construct prompt:", err)
     80 		}
     81 		msg := openai.Message{
     82 			Role: openai.RoleUser, Content: buf.String(),
     83 		}
     84 		chat.Messages = append(chat.Messages, msg)
     85 		reply, err := client.Complete(&chat)
     86 		if err != nil {
     87 			log.Fatalln("llm complete:", err)
     88 		}
     89 		fmt.Println(reply.Content)
     90 		return
     91 	}
     92 
     93 	sc := bufio.NewScanner(os.Stdin)
     94 	if len(flag.Args()) > 0 {
     95 		log.Println("conversation mode, ignoring arguments")
     96 	}
     97 	for sc.Scan() {
     98 		if sc.Text() == "." {
     99 			msg := openai.Message{
    100 				Role:    openai.RoleUser,
    101 				Content: buf.String(),
    102 			}
    103 			chat.Messages = append(chat.Messages, msg)
    104 			reply, err := client.Complete(&chat)
    105 			if err != nil {
    106 				fmt.Fprintln(os.Stderr, "chat not completed:", err)
    107 				continue // try again, allowing a retry with another "." line
    108 			}
    109 			buf.Reset()
    110 			fmt.Println(reply.Content)
    111 			chat.Messages = append(chat.Messages, *reply)
    112 			continue
    113 		}
    114 		fmt.Fprintln(buf, sc.Text())
    115 	}
    116 }