ant: improve encapsulation
- Replace string literals with package constants for message roles and content types.
- Create UserStringMessage helper function to simplify user message creation
- Replace manual Content creation with ant.StringContent()
Co-Authored-By: sketch <hello@sketch.dev>
diff --git a/loop/agent.go b/loop/agent.go
index adbdf51..315ff89 100644
--- a/loop/agent.go
+++ b/loop/agent.go
@@ -613,7 +613,7 @@
if resp.StopReason == ant.StopReasonToolUse {
var toolCalls []ToolCall
for _, part := range resp.Content {
- if part.Type == "tool_use" {
+ if part.Type == ant.ContentTypeToolUse {
toolCalls = append(toolCalls, ToolCall{
Name: part.ToolName,
Input: string(part.ToolInput),
@@ -997,13 +997,13 @@
case <-ctx.Done():
return m, ctx.Err()
case msg := <-a.inbox:
- m = append(m, ant.Content{Type: "text", Text: msg})
+ m = append(m, ant.StringContent(msg))
}
}
for {
select {
case msg := <-a.inbox:
- m = append(m, ant.Content{Type: "text", Text: msg})
+ m = append(m, ant.StringContent(msg))
default:
return m, nil
}
@@ -1087,7 +1087,7 @@
}
userMessage := ant.Message{
- Role: "user",
+ Role: ant.MessageRoleUser,
Content: msgs,
}
@@ -1211,19 +1211,19 @@
// Inject any auto-generated messages from quality checks
for _, msg := range autoqualityMessages {
- msgs = append(msgs, ant.Content{Type: "text", Text: msg})
+ msgs = append(msgs, ant.StringContent(msg))
}
// Handle cancellation by appending a message about it
if cancelled {
- msgs = append(msgs, ant.Content{Type: "text", Text: cancelToolUseMessage})
+ msgs = append(msgs, ant.StringContent(cancelToolUseMessage))
// EndOfTurn is false here so that the client of this agent keeps processing
// further messages; the conversation is not over.
a.pushToOutbox(ctx, AgentMessage{Type: ErrorMessageType, Content: userCancelMessage, EndOfTurn: false})
} else if err := a.convo.OverBudget(); err != nil {
// Handle budget issues by appending a message about it
budgetMsg := "We've exceeded our budget. Please ask the user to confirm before continuing by ending the turn."
- msgs = append(msgs, ant.Content{Type: "text", Text: budgetMsg})
+ msgs = append(msgs, ant.StringContent(budgetMsg))
a.pushToOutbox(ctx, budgetMessage(fmt.Errorf("warning: %w (ask to keep trying, if you'd like)", err)))
}
@@ -1233,7 +1233,7 @@
// Send the combined message to continue the conversation
a.stateMachine.Transition(ctx, StateSendingToolResults, "Sending tool results back to LLM")
resp, err := a.convo.SendMessage(ant.Message{
- Role: "user",
+ Role: ant.MessageRoleUser,
Content: results,
})
if err != nil {
@@ -1268,7 +1268,7 @@
// Collect all text content
var allText strings.Builder
for _, content := range msg.Content {
- if content.Type == "text" && content.Text != "" {
+ if content.Type == ant.ContentTypeText && content.Text != "" {
if allText.Len() > 0 {
allText.WriteString("\n\n")
}
@@ -1604,10 +1604,7 @@
Reply with ONLY the reprompt text.
`
- userMessage := ant.Message{
- Role: "user",
- Content: []ant.Content{{Type: "text", Text: msg}},
- }
+ userMessage := ant.UserStringMessage(msg)
// By doing this in a subconversation, the agent doesn't call tools (because
// there aren't any), and there's not a concurrency risk with on-going other
// outstanding conversations.
diff --git a/loop/agent_test.go b/loop/agent_test.go
index f1d5b51..24c1c55 100644
--- a/loop/agent_test.go
+++ b/loop/agent_test.go
@@ -492,7 +492,7 @@
}
func (m *mockConvoInterface) SendUserTextMessage(s string, otherContents ...ant.Content) (*ant.MessageResponse, error) {
- return m.SendMessage(ant.Message{Role: "user", Content: []ant.Content{{Type: "text", Text: s}}})
+ return m.SendMessage(ant.UserStringMessage(s))
}
func (m *mockConvoInterface) ToolResultContents(ctx context.Context, resp *ant.MessageResponse) ([]ant.Content, error) {
@@ -503,7 +503,7 @@
}
func (m *mockConvoInterface) ToolResultCancelContents(resp *ant.MessageResponse) ([]ant.Content, error) {
- return []ant.Content{{Type: "text", Text: "Tool use cancelled"}}, nil
+ return []ant.Content{ant.StringContent("Tool use cancelled")}, nil
}
func (m *mockConvoInterface) CancelToolUse(toolUseID string, cause error) error {
@@ -544,7 +544,7 @@
return &ant.MessageResponse{
StopReason: ant.StopReasonEndTurn,
Content: []ant.Content{
- {Type: "text", Text: "This is a test response"},
+ ant.StringContent("This is a test response"),
},
}, nil
}
@@ -619,8 +619,8 @@
return &ant.MessageResponse{
StopReason: ant.StopReasonToolUse,
Content: []ant.Content{
- {Type: "text", Text: "I'll use a tool"},
- {Type: "tool_use", ToolName: "test_tool", ToolInput: []byte("{}"), ID: "test_id"},
+ ant.StringContent("I'll use a tool"),
+ {Type: ant.ContentTypeToolUse, ToolName: "test_tool", ToolInput: []byte("{}"), ID: "test_id"},
},
}, nil
}
@@ -628,14 +628,14 @@
return &ant.MessageResponse{
StopReason: ant.StopReasonEndTurn,
Content: []ant.Content{
- {Type: "text", Text: "Finished using the tool"},
+ ant.StringContent("Finished using the tool"),
},
}, nil
}
// Tool result content handler
mockConvo.ToolResultContentsFunc = func(ctx context.Context, resp *ant.MessageResponse) ([]ant.Content, error) {
- return []ant.Content{{Type: "text", Text: "Tool executed successfully"}}, nil
+ return []ant.Content{ant.StringContent("Tool executed successfully")}, nil
}
// Track state transitions
diff --git a/loop/agent_user_cancel_test.go b/loop/agent_user_cancel_test.go
index f79e73e..e55ba17 100644
--- a/loop/agent_user_cancel_test.go
+++ b/loop/agent_user_cancel_test.go
@@ -21,12 +21,7 @@
inbox: make(chan string, 1),
outbox: make(chan AgentMessage, 1),
}
- userMsg := ant.Message{
- Role: "user",
- Content: []ant.Content{
- {Type: "text", Text: "hi"},
- },
- }
+ userMsg := ant.UserStringMessage("hi")
userMsgResponse := &ant.MessageResponse{}
mockConvo.ExpectCall("SendMessage", userMsg).Return(userMsgResponse, nil)
@@ -55,9 +50,9 @@
outbox: make(chan AgentMessage, 1),
}
userMsg := ant.Message{
- Role: "user",
+ Role: ant.MessageRoleUser,
Content: []ant.Content{
- {Type: "text", Text: "hi"},
+ {Type: ant.ContentTypeText, Text: "hi"},
},
}
userMsgResponse := &ant.MessageResponse{
@@ -86,7 +81,7 @@
},
}
toolUseResultsMsg := ant.Message{
- Role: "user",
+ Role: ant.MessageRoleUser,
Content: toolUseContents,
}
toolUseResponse := &ant.MessageResponse{
@@ -132,12 +127,7 @@
inbox: make(chan string, 1),
outbox: make(chan AgentMessage, 10), // don't let anything block on outbox.
}
- userMsg := ant.Message{
- Role: "user",
- Content: []ant.Content{
- {Type: "text", Text: "hi"},
- },
- }
+ userMsg := ant.UserStringMessage("hi")
userMsgResponse := &ant.MessageResponse{
StopReason: ant.StopReasonToolUse,
Content: []ant.Content{
@@ -153,12 +143,7 @@
OutputTokens: 200,
},
}
- toolUseResultsMsg := ant.Message{
- Role: "user",
- Content: []ant.Content{
- {Type: "text", Text: cancelToolUseMessage},
- },
- }
+ toolUseResultsMsg := ant.UserStringMessage(cancelToolUseMessage)
toolUseResponse := &ant.MessageResponse{
StopReason: ant.StopReasonEndTurn,
Content: []ant.Content{
@@ -223,9 +208,9 @@
outbox: make(chan AgentMessage, 10), // don't let anything block on outbox.
}
userMsg := ant.Message{
- Role: "user",
+ Role: ant.MessageRoleUser,
Content: []ant.Content{
- {Type: "text", Text: "hi"},
+ {Type: ant.ContentTypeText, Text: "hi"},
},
}
userMsgResponse := &ant.MessageResponse{
@@ -244,9 +229,9 @@
},
}
toolUseResultsMsg := ant.Message{
- Role: "user",
+ Role: ant.MessageRoleUser,
Content: []ant.Content{
- {Type: "text", Text: cancelToolUseMessage},
+ {Type: ant.ContentTypeText, Text: cancelToolUseMessage},
},
}
toolUseResultResponse := &ant.MessageResponse{
@@ -263,9 +248,9 @@
},
}
userFollowUpMsg := ant.Message{
- Role: "user",
+ Role: ant.MessageRoleUser,
Content: []ant.Content{
- {Type: "text", Text: "that was the wrong thing to do"},
+ {Type: ant.ContentTypeText, Text: "that was the wrong thing to do"},
},
}
userFollowUpResponse := &ant.MessageResponse{
@@ -339,12 +324,7 @@
// Define test message
// This simulates something that would result in claude responding with tool_use responses.
- userMsg := ant.Message{
- Role: "user",
- Content: []ant.Content{
- {Type: "text", Text: "use test_tool for something"},
- },
- }
+ userMsg := ant.UserStringMessage("use test_tool for something")
// Mock initial response with tool use
userMsgResponse := &ant.MessageResponse{
StopReason: ant.StopReasonToolUse,
@@ -370,11 +350,8 @@
},
}
canceledToolUseMsg := ant.Message{
- Role: "user",
- Content: append(canceledToolUseContents, ant.Content{
- Type: ant.ContentTypeText,
- Text: cancelToolUseMessage,
- }),
+ Role: ant.MessageRoleUser,
+ Content: append(canceledToolUseContents, ant.StringContent(cancelToolUseMessage)),
}
// Set up expected behaviors
waitForSendMessage := make(chan any)