trim images

server/images.go

@@ -58,11 +58,17 @@ type Message struct {
 	Content string `json:"content"`
 }
 
+type ImageData struct {
+	Rank int
+	api.ImageData
+}
+
 type PromptVars struct {
 	System   string
 	Prompt   string
 	Response string
 	First    bool
+	Images   []ImageData
 }
 
 // extractParts extracts the parts of the template before and after the {{.Response}} node.
@@ -147,15 +153,13 @@ func (m *Model) PostResponseTemplate(p PromptVars) (string, error) {
 }
 
 type ChatHistory struct {
-	Prompts       []PromptVars
-	CurrentImages []api.ImageData
-	LastSystem    string
+	Prompts    []PromptVars
+	LastSystem string
 }
 
 // ChatPrompts returns a list of formatted chat prompts from a list of messages
 func (m *Model) ChatPrompts(msgs []api.Message) (*ChatHistory, error) {
 	// build the prompt from the list of messages
-	var currentImages []api.ImageData
 	lastSystem := m.System
 	currentVars := PromptVars{
 		First:  true,
@@ -163,6 +167,7 @@ func (m *Model) ChatPrompts(msgs []api.Message) (*ChatHistory, error) {
 	}
 
 	prompts := []PromptVars{}
+	var images []ImageData
 
 	for _, msg := range msgs {
 		switch strings.ToLower(msg.Role) {
@@ -182,10 +187,15 @@ func (m *Model) ChatPrompts(msgs []api.Message) (*ChatHistory, error) {
 
 			currentVars.Prompt = msg.Content
 			for i := range msg.Images {
-				currentVars.Prompt += fmt.Sprintf(" [img-%d]", len(currentImages)+i)
+				currentVars.Prompt += fmt.Sprintf(" [img-%d]", len(images)+i)
+				currentVars.Images = append(currentVars.Images, ImageData{
+					Rank:      len(images) + i,
+					ImageData: msg.Images[i],
+				})
+
 			}
 
-			currentImages = append(currentImages, msg.Images...)
+			images = append(images, currentVars.Images...)
 		case "assistant":
 			currentVars.Response = msg.Content
 			prompts = append(prompts, currentVars)
@@ -201,9 +211,8 @@ func (m *Model) ChatPrompts(msgs []api.Message) (*ChatHistory, error) {
 	}
 
 	return &ChatHistory{
-		Prompts:       prompts,
-		CurrentImages: currentImages,
-		LastSystem:    lastSystem,
+		Prompts:    prompts,
+		LastSystem: lastSystem,
 	}, nil
 }