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) {

		input := ""
		if len(args) == 0 {
			stdin, err := io.ReadAll(os.Stdin)
			if err != nil {
				log.Fatalf("failed to read input: %q", err)
				return
			}

			input = string(stdin)
		} else {
			input = strings.Join(args, " ")
		}

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

		lines := strings.Split(string(input), "\n")
		if strings.HasPrefix(lines[0], "#") {
			temperature = getTemperature(lines[0])
			maxTokens = getMaxTokens(lines[0])
			lines = lines[0:]
		} else {
			if m, err := f.GetInt("max_tokens"); err == nil {
				maxTokens = m
			}

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

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

		var isFirst = false
		if len(prompts) == 1 && prompts[0].Type != chat.MessageTypeSystem {
			isFirst = true
			prf, _ := f.GetString("system-prompt")
			promptStr, err := os.ReadFile(prf)
			if err != nil {
				promptStr = []byte(DefaultSystemPrompt)
			}

			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
			}
		}

		messages := chat.New(prompts...)
		resp, err := llm.Chat(context.TODO(), messages)
		if err != nil {
			log.Fatalf("failed to send send: %q", err.Error())
			return
		}

		term, _ := f.GetBool("term")

		if isatty.IsTerminal(os.Stdout.Fd()) || isatty.IsCygwinTerminal(os.Stderr.Fd()) {
			if term {
				resp = string(markdown.Render(resp, 90, 4))

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

				tmpf, err := os.CreateTemp(os.TempDir(), "*.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)
				for _, p := range prompts {
					prefix := ""
					switch p.Type {
					case chat.MessageTypeSystem:
						prefix = "SYSTEM: "
					case chat.MessageTypeUser:
						prefix = "USER: "
					case chat.MessageTypeAssistant:
						prefix = "AI: "
					}

					fmt.Fprintf(tmpf, "%s%s\n\n", prefix, strings.Trim(p.Prompt.String(), "\r\n"))
				}

				fmt.Fprintf(tmpf, "AI: %s\n\nUSER: ", resp)
				tmpf.Close()

				cmd := exec.Command("vim", "-c", "norm! 4j", tmpf.Name())
				cmd.Stdin = os.Stdin
				cmd.Stdout = os.Stdout
				cmd.Stderr = os.Stderr

				cmd.Run()
			}
		} else {

			fmt.Println(resp)
		}
	},
}

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 := io.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(string(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 := io.ReadAll(os.Stdin)
		if err != nil {
			log.Fatalf("failed to read input: %q", err)
		}

		toks, err := vichat.Tokenize(string(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"`
}