diff --git a/pkg/cmd/util/util.go b/pkg/cmd/util/util.go index 7c4b962..84d3178 100644 --- a/pkg/cmd/util/util.go +++ b/pkg/cmd/util/util.go @@ -2,6 +2,7 @@ package util import ( "context" + "encoding/json" "fmt" "os" "strings" @@ -149,36 +150,40 @@ func FormatForExternalPrompt(messages []model.Message, system bool) string { } func GenerateTitle(ctx *lmcli.Context, messages []model.Message) (string, error) { - const prompt = `Above is an excerpt from a conversation between a user and AI assistant. - -Reply with a short title (no more than 8 words) that reflects the topic of the conversation, read from the user's perspective. + 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. Example conversation: -""" -User: +[{"role": "user", "content": "Can you help me with my math homework?"},{"role": "assistant", "content": "Sure, what topic are you struggling with?"}] - Hello! +Your response (response only with JSON in the format shown): -Assistant: - - Hello! How may I assist you? -""" - -Your response: - -""" -Title: A brief introduction -""" - -Respond only as shown above. +{"title": "Help with math homework"} ` - conversation := FormatForExternalPrompt(messages, false) + type msg struct { + Role string + Content string + } + + var msgs []msg + for _, m := range messages { + 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 := []model.Message{ + { + Role: model.MessageRoleSystem, + Content: systemPrompt, + }, { Role: model.MessageRoleUser, - Content: fmt.Sprintf("\"\"\"\n%s\n\"\"\"\n\n%s", conversation, prompt), + Content: string(conversation), }, } @@ -197,12 +202,16 @@ Respond only as shown above. return "", err } - // TODO: validate response format - response = strings.TrimPrefix(response, "Title:") - response = strings.TrimSpace(response) - response = strings.Trim(response, "\"") + // Parse the JSON response + var jsonResponse struct { + Title string `json:"title"` + } + err = json.Unmarshal([]byte(response), &jsonResponse) + if err != nil { + return "", err + } - return response, nil + return jsonResponse.Title, nil } // ShowWaitAnimation prints an animated ellipses to stdout until something is