6d95ecc · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix: skip image attachments in history when model doesn't support them (#2818)

6d95eccAndrey Nering5 months ago

internal/agent/agent.go+26−4

compare with main

unifiedsplit

@@ -245,7 +245,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

245 defer cancel()

246 defer a.activeRequests.Del(call.SessionID)

247

248− history, files := a.preparePrompt(msgs, call.Attachments...)

248+ history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...)

249

250 startTime := time.Now()

251 a.eventPromptSent(call.SessionID)

@@ -643,7 +643,7 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan

643 return nil

644 }

645

646− aiMsgs, _ := a.preparePrompt(msgs)

646+ aiMsgs, _ := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages)

647

648 genCtx, cancel := context.WithCancel(ctx)

649 a.activeRequests.Set(sessionID, cancel)

@@ -786,7 +786,7 @@ func (a *sessionAgent) createUserMessage(ctx context.Context, call SessionAgentC

786 return msg, nil

787}

788

789−func (a *sessionAgent) preparePrompt(msgs []message.Message, attachments ...message.Attachment) ([]fantasy.Message, []fantasy.FilePart) {

789+func (a *sessionAgent) preparePrompt(msgs []message.Message, supportsImages bool, attachments ...message.Attachment) ([]fantasy.Message, []fantasy.FilePart) {

790 var history []fantasy.Message

791 if !a.isSubAgent {

792 history = append(history, fantasy.NewUserMessage(

@@ -830,7 +830,15 @@ If not, please feel free to ignore. Again do not mention this message to the use

830 }

831 continue

832 }

833+ aiMsgs := m.ToAIMessage()

834+ if !supportsImages {

835+ for i := range aiMsgs {

836+ if aiMsgs[i].Role == fantasy.MessageRoleUser {

837+ aiMsgs[i].Content = filterFileParts(aiMsgs[i].Content)

838+ }

839+ }

840+ }

833− history = append(history, m.ToAIMessage()...)

841+ history = append(history, aiMsgs...)

834

842

835 if m.Role == message.Assistant {

843 if m.Role == message.Assistant {

836 if msg, ok := syntheticToolResultsForOrphanedCalls(m, knownToolResultIDs); ok {

844 if msg, ok := syntheticToolResultsForOrphanedCalls(m, knownToolResultIDs); ok {

@@ -854,6 +862,20 @@ If not, please feel free to ignore. Again do not mention this message to the use

854 return history, files

862 return history, files

855}

863}

856

864

865+// filterFileParts removes fantasy.FilePart entries from a slice of message

866+// parts. Used to strip image attachments from historical user messages when

867+// the current model does not support them.

868+func filterFileParts(parts []fantasy.MessagePart) []fantasy.MessagePart {

869+ filtered := make([]fantasy.MessagePart, 0, len(parts))

870+ for _, part := range parts {

871+ if _, ok := fantasy.AsMessagePart[fantasy.FilePart](part); ok {

872+ continue

873+ }

874+ filtered = append(filtered, part)

875+ }

876+ return filtered

877+}

878+

857// filterOrphanedToolResults converts a tool message to a fantasy.Message,

879// filterOrphanedToolResults converts a tool message to a fantasy.Message,

858// dropping any tool result parts whose tool_call_id has no matching tool call

880// dropping any tool result parts whose tool_call_id has no matching tool call

859// in the known set. An orphaned result causes API validation to fail on every

881// in the known set. An orphaned result causes API validation to fail on every

internal/agent/agent_test.go+47−2

@@ -656,6 +656,51 @@ func BenchmarkBuildSummaryPrompt(b *testing.B) {

656 }

657}

658

659+func TestPreparePrompt_FiltersImageAttachments(t *testing.T) {

660+ env := testEnv(t)

661+ sa := testSessionAgent(env, nil, nil, "test prompt")

662+ agent := sa.(*sessionAgent)

663+

664+ ctx := t.Context()

665+ sess, err := env.sessions.Create(ctx, "test")

666+ require.NoError(t, err)

667+

668+ // User message with text, a text attachment, and an image attachment.

669+ _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{

670+ Role: message.User,

671+ Parts: []message.ContentPart{

672+ message.TextContent{Text: "hello world"},

673+ message.BinaryContent{Path: "notes.txt", MIMEType: "text/plain", Data: []byte("important notes")},

674+ message.BinaryContent{Path: "image.png", MIMEType: "image/png", Data: []byte("fake-image-data")},

675+ },

676+ })

677+ require.NoError(t, err)

678+

679+ msgs, err := env.messages.List(ctx, sess.ID)

680+ require.NoError(t, err)

681+

682+ // When supportsImages is false, image attachments should be stripped.

683+ history, _ := agent.preparePrompt(msgs, false)

684+ // First message is the system reminder, second is the user message.

685+ require.Len(t, history, 2)

686+ require.Len(t, history[1].Content, 1)

687+ text, ok := fantasy.AsMessagePart[fantasy.TextPart](history[1].Content[0])

688+ require.True(t, ok)

689+ require.Contains(t, text.Text, "hello world")

690+ require.Contains(t, text.Text, "important notes")

691+

692+ // When supportsImages is true, image attachments should remain.

693+ history, _ = agent.preparePrompt(msgs, true)

694+ require.Len(t, history, 2)

695+ require.Len(t, history[1].Content, 2)

696+ text, ok = fantasy.AsMessagePart[fantasy.TextPart](history[1].Content[0])

697+ require.True(t, ok)

698+ require.Contains(t, text.Text, "hello world")

699+ file, ok := fantasy.AsMessagePart[fantasy.FilePart](history[1].Content[1])

700+ require.True(t, ok)

701+ require.Equal(t, "image.png", file.Filename)

702+}

703+

659func TestPreparePrompt_OrphanedToolUse(t *testing.T) {

704func TestPreparePrompt_OrphanedToolUse(t *testing.T) {

660 env := testEnv(t)

705 env := testEnv(t)

661 sa := testSessionAgent(env, nil, nil, "test prompt")

706 sa := testSessionAgent(env, nil, nil, "test prompt")

@@ -702,7 +747,7 @@ func TestPreparePrompt_OrphanedToolUse(t *testing.T) {

702 msgs, err := env.messages.List(ctx, sess.ID)

747 msgs, err := env.messages.List(ctx, sess.ID)

703 require.NoError(t, err)

748 require.NoError(t, err)

704

749

705− history, _ := agent.preparePrompt(msgs)

750+ history, _ := agent.preparePrompt(msgs, true)

706

751

707 // The history must contain a synthetic tool result for the orphaned call.

752 // The history must contain a synthetic tool result for the orphaned call.

708 found := false

753 found := false

@@ -776,7 +821,7 @@ func TestPreparePrompt_OrphanedToolUseMixed(t *testing.T) {

776 msgs, err := env.messages.List(ctx, sess.ID)

821 msgs, err := env.messages.List(ctx, sess.ID)

777 require.NoError(t, err)

822 require.NoError(t, err)

778

823

779− history, _ := agent.preparePrompt(msgs)

824+ history, _ := agent.preparePrompt(msgs, true)

780

825

781 // Should have a synthetic result only for the orphaned call.

826 // Should have a synthetic result only for the orphaned call.

782 var syntheticCount int

827 var syntheticCount int