From af613bab335bc6d62e059d200c4b7fec0ade1817 Mon Sep 17 00:00:00 2001
From: Patrick Devine <patrick@infrahq.com>
Date: Wed, 9 Oct 2024 14:24:36 -0700
Subject: [PATCH] fix prompting

---
 server/prompt.go      |  2 +-
 server/prompt_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/server/prompt.go b/server/prompt.go
index eec7ce15..d0cc8e9d 100644
--- a/server/prompt.go
+++ b/server/prompt.go
@@ -77,7 +77,7 @@ func chatPrompt(ctx context.Context, m *Model, tokenize tokenizeFunc, opts *api.
 
 	if isMllama {
 		lastMsgIdx := len(msgs) - 1
-		for i := lastMsgIdx; i > currMsgIdx; i-- {
+		for i := lastMsgIdx; i >= currMsgIdx; i-- {
 			if len(msgs[i].Images) > 0 {
 				data, aspectRatioID, err := imageproc.Preprocess(msgs[i].Images[0])
 				if err != nil {
diff --git a/server/prompt_test.go b/server/prompt_test.go
index d02a6e93..d437ebb4 100644
--- a/server/prompt_test.go
+++ b/server/prompt_test.go
@@ -229,6 +229,19 @@ func TestChatPrompt(t *testing.T) {
 				prompt: "You're a test, Harry! I-I'm a what? A test. And a thumping good one at that, I'd wager. ",
 			},
 		},
+		{
+			name:  "messages with mllama single prompt",
+			model: mllamaModel,
+			limit: 2048,
+			msgs: []api.Message{
+				{Role: "user", Content: "How many hotdogs are in this image?", Images: []api.ImageData{imgBuf}},
+			},
+			expect: expect{
+				prompt:        "<|image|>How many hotdogs are in this image? ",
+				images:        [][]byte{imgBuf},
+				aspectRatioID: 1,
+			},
+		},
 		{
 			name:  "messages with mllama",
 			model: mllamaModel,
@@ -259,6 +272,21 @@ func TestChatPrompt(t *testing.T) {
 				aspectRatioID: 1,
 			},
 		},
+		{
+			name:  "earlier image with mllama",
+			model: mllamaModel,
+			limit: 2048,
+			msgs: []api.Message{
+				{Role: "user", Content: "How many hotdogs are in this image?", Images: []api.ImageData{imgBuf}},
+				{Role: "assistant", Content: "There are four hotdogs."},
+				{Role: "user", Content: "Which ones have mustard?"},
+			},
+			expect: expect{
+				prompt:        "<|image|>How many hotdogs are in this image? There are four hotdogs. Which ones have mustard? ",
+				images:        [][]byte{imgBuf},
+				aspectRatioID: 1,
+			},
+		},
 		{
 			name:  "too many images with mllama",
 			model: mllamaModel,