package util import ( "context" "encoding/json" "fmt" "os" "strings" "time" "git.mlow.ca/mlow/lmcli/pkg/api" "git.mlow.ca/mlow/lmcli/pkg/lmcli" "git.mlow.ca/mlow/lmcli/pkg/util" "github.com/charmbracelet/lipgloss" ) // Prompt prompts the configured the configured model and streams the response // to stdout. Returns all model reply messages. func Prompt(ctx *lmcli.Context, messages []api.Message, callback func(api.Message)) (*api.Message, error) { m, provider, err := ctx.GetModelProvider(*ctx.Config.Defaults.Model) if err != nil { return nil, err } params := api.RequestParameters{ Model: m, MaxTokens: *ctx.Config.Defaults.MaxTokens, Temperature: *ctx.Config.Defaults.Temperature, } system := ctx.DefaultSystemPrompt() agent := ctx.GetAgent(ctx.Config.Defaults.Agent) if agent != nil { if agent.SystemPrompt != "" { system = agent.SystemPrompt } params.Toolbox = agent.Toolbox } if system != "" { messages = api.ApplySystemPrompt(messages, system, false) } content := make(chan api.Chunk) defer close(content) // render the content received over the channel go ShowDelayedContent(content) reply, err := provider.CreateChatCompletionStream( context.Background(), params, messages, content, ) if reply.Content != "" { // there was some content, so break to a new line after it fmt.Println() if err != nil { lmcli.Warn("Received partial response. Error: %v\n", err) err = nil } } return reply, err } // lookupConversation either returns the conversation found by the // short name or exits the program func LookupConversation(ctx *lmcli.Context, shortName string) *api.Conversation { c, err := ctx.Store.ConversationByShortName(shortName) if err != nil { lmcli.Fatal("Could not lookup conversation: %v\n", err) } if c.ID == 0 { lmcli.Fatal("Conversation not found: %s\n", shortName) } return c } func LookupConversationE(ctx *lmcli.Context, shortName string) (*api.Conversation, error) { c, err := ctx.Store.ConversationByShortName(shortName) if err != nil { return nil, fmt.Errorf("Could not lookup conversation: %v", err) } if c.ID == 0 { return nil, fmt.Errorf("Conversation not found: %s", shortName) } return c, nil } func HandleConversationReply(ctx *lmcli.Context, c *api.Conversation, persist bool, toSend ...api.Message) { messages, err := ctx.Store.PathToLeaf(c.SelectedRoot) if err != nil { lmcli.Fatal("Could not load messages: %v\n", err) } HandleReply(ctx, &messages[len(messages)-1], persist, toSend...) } // handleConversationReply handles sending messages to an existing // conversation, optionally persisting both the sent replies and responses. func HandleReply(ctx *lmcli.Context, to *api.Message, persist bool, messages ...api.Message) { if to == nil { lmcli.Fatal("Can't prompt from an empty message.") } existing, err := ctx.Store.PathToRoot(to) if err != nil { lmcli.Fatal("Could not load messages: %v\n", err) } RenderConversation(ctx, append(existing, messages...), true) var savedReplies []api.Message if persist && len(messages) > 0 { savedReplies, err = ctx.Store.Reply(to, messages...) if err != nil { lmcli.Warn("Could not save messages: %v\n", err) } } // render a message header with no contents RenderMessage(ctx, (&api.Message{Role: api.MessageRoleAssistant})) var lastSavedMessage *api.Message lastSavedMessage = to if len(savedReplies) > 0 { lastSavedMessage = &savedReplies[len(savedReplies)-1] } replyCallback := func(reply api.Message) { if !persist { return } savedReplies, err = ctx.Store.Reply(lastSavedMessage, reply) if err != nil { lmcli.Warn("Could not save reply: %v\n", err) } lastSavedMessage = &savedReplies[0] } _, err = Prompt(ctx, append(existing, messages...), replyCallback) if err != nil { lmcli.Fatal("Error fetching LLM response: %v\n", err) } } func FormatForExternalPrompt(messages []api.Message, system bool) string { sb := strings.Builder{} for _, message := range messages { if message.Content == "" { continue } switch message.Role { case api.MessageRoleAssistant, api.MessageRoleToolCall: sb.WriteString("Assistant:\n\n") case api.MessageRoleUser: sb.WriteString("User:\n\n") default: continue } sb.WriteString(fmt.Sprintf("%s", lipgloss.NewStyle().PaddingLeft(1).Render(message.Content))) } return sb.String() } func GenerateTitle(ctx *lmcli.Context, messages []api.Message) (string, error) { const systemPrompt = `You will be shown a conversation between a user and an AI assistant. Your task is to generate a short title (8 words or less) for the provided conversation that reflects the conversation's topic. Your response is expected to be in JSON in the format shown below. Example conversation: [{"role": "user", "content": "Can you help me with my math homework?"},{"role": "assistant", "content": "Sure, what topic are you struggling with?"}] Example response: {"title": "Help with math homework"} ` type msg struct { Role string Content string } var msgs []msg for _, m := range messages { switch m.Role { case api.MessageRoleAssistant, api.MessageRoleUser: msgs = append(msgs, msg{string(m.Role), m.Content}) } } // Serialize the conversation to JSON conversation, err := json.Marshal(msgs) if err != nil { return "", err } generateRequest := []api.Message{ { Role: api.MessageRoleSystem, Content: systemPrompt, }, { Role: api.MessageRoleUser, Content: string(conversation), }, } m, provider, err := ctx.GetModelProvider( *ctx.Config.Conversations.TitleGenerationModel, ) if err != nil { return "", err } requestParams := api.RequestParameters{ Model: m, MaxTokens: 25, } response, err := provider.CreateChatCompletion( context.Background(), requestParams, generateRequest, ) if err != nil { return "", err } // Parse the JSON response var jsonResponse struct { Title string `json:"title"` } err = json.Unmarshal([]byte(response.Content), &jsonResponse) if err != nil { return "", err } return jsonResponse.Title, nil } // ShowWaitAnimation prints an animated ellipses to stdout until something is // received on the signal channel. An empty string sent to the channel to // notify the caller that the animation has completed (carriage returned). func ShowWaitAnimation(signal chan any) { // Save the current cursor position fmt.Print("\033[s") animationStep := 0 for { select { case _ = <-signal: // Relmcli the cursor position fmt.Print("\033[u") signal <- "" return default: // Move the cursor to the saved position modSix := animationStep % 6 if modSix == 3 || modSix == 0 { fmt.Print("\033[u") } if modSix < 3 { fmt.Print(".") } else { fmt.Print(" ") } animationStep++ time.Sleep(250 * time.Millisecond) } } } // ShowDelayedContent displays a waiting animation to stdout while waiting // for content to be received on the provided channel. As soon as any (possibly // chunked) content is received on the channel, the waiting animation is // replaced by the content. // Blocks until the channel is closed. func ShowDelayedContent(content <-chan api.Chunk) { waitSignal := make(chan any) go ShowWaitAnimation(waitSignal) firstChunk := true for chunk := range content { if firstChunk { // notify wait animation that we've received data waitSignal <- "" // wait for signal that wait animation has completed <-waitSignal firstChunk = false } fmt.Print(chunk.Content) } } // RenderConversation renders the given messages to TTY, with optional space // for a subsequent message. spaceForResponse controls how many '\n' characters // are printed immediately after the final message (1 if false, 2 if true) func RenderConversation(ctx *lmcli.Context, messages []api.Message, spaceForResponse bool) { l := len(messages) for i, message := range messages { RenderMessage(ctx, &message) if i < l-1 || spaceForResponse { // print an additional space before the next message fmt.Println() } } } func RenderMessage(ctx *lmcli.Context, m *api.Message) { var messageAge string if m.CreatedAt.IsZero() { messageAge = "now" } else { now := time.Now() messageAge = util.HumanTimeElapsedSince(now.Sub(m.CreatedAt)) } headerStyle := lipgloss.NewStyle().Bold(true) switch m.Role { case api.MessageRoleSystem: headerStyle = headerStyle.Foreground(lipgloss.Color("9")) // bright red case api.MessageRoleUser: headerStyle = headerStyle.Foreground(lipgloss.Color("10")) // bright green case api.MessageRoleAssistant: headerStyle = headerStyle.Foreground(lipgloss.Color("12")) // bright blue } role := headerStyle.Render(m.Role.FriendlyRole()) separatorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("3")) separator := separatorStyle.Render("===") timestamp := separatorStyle.Render(messageAge) fmt.Printf("%s %s - %s %s\n\n", separator, role, timestamp, separator) if m.Content != "" { ctx.Chroma.Highlight(os.Stdout, m.Content) fmt.Println() } }