cmd

Documentation

Index

• Constants
• Variables
• func CreatePrompts(lines []string) []chat.PromptMessage
• type TimeQuery
• type TimeResp

Constants

const DefaultMaxTokens = 1000

const DefaultSystemPrompt = "You are a helpful assistant."

const DefaultTemperature = 0.7

Variables

var ChatCmd = &cobra.Command{
	Use:   "chat",
	Short: "read a chat from stdin and send to LLM chat",
	Run: func(cmd *cobra.Command, args []string) {

		if hasDifference(vimPlugins, "vim/ftdetect", os.ExpandEnv("$HOME/.vim/ftdetect")) {
			log.Printf("WARNING: your vichat vim plugin appears to be out of sync. run `vichat i` to install it again.")
		}

		if hasDifference(vimPlugins, "vim/ftplugin", os.ExpandEnv("$HOME/.vim/ftplugin")) {
			log.Printf("WARNING: your vichat vim plugin appears to be out of sync. run `vichat i` to install it again.")
		}

		if hasDifference(vimPlugins, "vim/syntax", os.ExpandEnv("$HOME/.vim/syntax")) {
			log.Printf("WARNING: your vichat vim plugin appears to be out of sync. run `vichat i` to install it again.")
		}

		var f = cmd.Flags()
		var temperature float32 = DefaultTemperature
		var maxTokens int = DefaultMaxTokens

		if m, err := f.GetInt("max_tokens"); err == nil {
			maxTokens = m
		}

		if t, err := f.GetFloat32("temperature"); err == nil {
			temperature = t
		}

		var input string
		var lines []string
		var isSimpleChat bool = false
		if !isatty.IsTerminal(os.Stdin.Fd()) {
			stdin, err := io.ReadAll(os.Stdin)
			if err != nil {
				log.Fatalf("failed to read input: %q", err)
				return
			}

			input = string(stdin)
			lines = strings.Split(string(input), "\n")
			if strings.HasPrefix(lines[0], "#") {
				temperature = getTemperature(lines[0])
				maxTokens = getMaxTokens(lines[0])
				lines = lines[0:]
			}
		} else {
			input = strings.Join(args, " ")
			lines = []string{input}
			isSimpleChat = true
		}

		llm := vichat.New().WithTemperature(temperature).WithMaxTokens(maxTokens)
		prompts := CreatePrompts(lines)
		if len(prompts) == 0 {
			log.Fatalf("invalid input")
			return
		}

		if isSimpleChat {
			var promptStr []byte
			var err error
			prf, _ := f.GetString("system-prompt")
			if prf == "system.prompt" {
				promptStr = []byte(DefaultSystemPrompt)
			} else {
				promptStr, err = os.ReadFile(prf)
				if err != nil {
					prd := csv.NewReader(bytes.NewReader(awesomePrompts))
					embedPrompts, err := prd.ReadAll()
					if err == nil {
						index := make([]string, len(embedPrompts))
						for i := range embedPrompts {
							index[i] = strings.ToLower(embedPrompts[i][0])
						}

						matches := fuzzy.RankFind(prf, index)
						sort.Sort(matches)

						hit := matches[0].OriginalIndex
						promptStr = []byte(embedPrompts[hit][1])
					}
				}
			}

			prompts = append([]chat.PromptMessage{{
				Type:   chat.MessageTypeSystem,
				Prompt: prompt.New(string(promptStr)),
			}}, prompts...)
		}

		if ok, _ := f.GetBool("func"); ok {
			if err := llm.BindFunction(
				getRelativeTime,
				"getRelativeTime",
				`Use this function to find out what time is it using a relative duration of seconds. 
				Translate the time into a num of seconds before calling the function. 
				e.g. 1 hour ago = getRelativeTime(3600)
					 now = getRelativeTime(0)
				`,
			); err != nil {
				log.Fatalf("failed to bind function: %q", err.Error())
				return
			}
		}

		isTermOutput, _ := f.GetBool("term")
		messages := chat.New(prompts...)
		if isatty.IsTerminal(os.Stdout.Fd()) {
			if isTermOutput {
				resp, err := llm.Chat(context.TODO(), messages)
				if err != nil {
					log.Fatalf("failed to send chat: %q", err.Error())
					return
				}

				resp = string(markdown.Render(resp, 90, 4))

				fmt.Println()
				fmt.Println(resp)
				fmt.Println()
			} else if isSimpleChat {

				dir, err := f.GetString("outdir")
				if err != nil {
					dir = os.TempDir()
				}

				tmpf, err := os.CreateTemp(dir, "*.chat")
				if err != nil {
					log.Fatalf("failed to create temp file: %q", err)
				}

				fmt.Fprintf(tmpf, "# temperature=%.1f, max_tokens=%d\n\n", temperature, maxTokens)
				printPrompts(tmpf, prompts)
				tmpf.Close()

				// invoke vim using cmd and open tmpf
				var cmd *exec.Cmd
				if input == "" {
					cmd = exec.Command("vim", "-c", "norm! GkA", tmpf.Name())
				} else {
					cmd = exec.Command("vim", "-c", "redraw|Chat", tmpf.Name())
				}

				cmd.Stdin = os.Stdin
				cmd.Stdout = os.Stdout

				cmd.Run()
			}
		} else {

			fmt.Printf("AI: ")

			err := llm.ChatStream(context.TODO(), func(s string) {
				fmt.Print(s)
			}, messages)

			if err != nil {
				log.Fatalf("failed to stream chat: %q", err.Error())
				return
			}

			fmt.Println("\n\nUSER: ")
		}
	},
}

var InstallCmd = &cobra.Command{
	Use:     "install-vim-plugin",
	Short:   "install the vim plugin",
	Aliases: []string{"i"},
	Run: func(cmd *cobra.Command, args []string) {
		installVim()
	},
}

var SplitCmd = &cobra.Command{
	Use:   "split",
	Short: "Split a text into multiple chunks",
	Run: func(cmd *cobra.Command, args []string) {
		text, err := readAll(os.Stdin)
		if err != nil {
			log.Fatalf("failed to read input: %q", err)
		}
		chunkSize, _ := cmd.Flags().GetInt("chunk-size")
		overlap, _ := cmd.Flags().GetInt("overlap")

		chunks := vichat.RecursiveTextSplit(text, chunkSize, overlap)
		for i := range chunks {
			fmt.Println("--------------------------------------------")
			fmt.Println(chunks[i])
		}
	},
}

var TokCmd = &cobra.Command{
	Use:   "tok",
	Short: "given a piece of text, tok estimate the num of tokens for a given model offline",
	Run: func(cmd *cobra.Command, args []string) {

		f := cmd.Flags()
		model, err := f.GetString("model")
		if err != nil {
			log.Fatalf("failed to read model: %q", err)
		}

		text, err := readAll(os.Stdin)
		if err != nil {
			log.Fatalf("failed to read input: %q", err)
		}

		toks, err := vichat.Tokenize(text, model)
		if err != nil {
			log.Fatalf("failed to tokenize: %q", err)
			return
		}

		fmt.Println(len(toks))
	},
}

Functions

func CreatePrompts

func CreatePrompts(lines []string) []chat.PromptMessage

Types

type TimeQuery

type TimeQuery struct {
	SecondsAgo int `json:"secondsAgo"`
}

type TimeResp

type TimeResp struct {
	Time  string `json:"time"`
	Query string `json:"query"`
}