diff --git a/pkg/agent/integration_test.go b/pkg/agent/integration_test.go index 2a2f052..68532bf 100644 --- a/pkg/agent/integration_test.go +++ b/pkg/agent/integration_test.go @@ -103,7 +103,7 @@ func TestIntegration_AgentLoop_Generate(t *testing.T) { MaxIters: 3, }) - resp, err := loop.Run(context.Background(), "What is 20+22? Use the add_numbers tool.") + resp, err := loop.Run(context.Background(), llm.Message{Role: llm.RoleUser, Content: "What is 20+22? Use the add_numbers tool."}) if err != nil { t.Fatalf("run failed: %v", err) } @@ -130,7 +130,7 @@ func TestIntegration_AgentLoop_Stream(t *testing.T) { MaxIters: 3, }) - stream := loop.RunStream(context.Background(), "Say something interesting.") + stream := loop.RunStream(context.Background(), llm.Message{Role: llm.RoleUser, Content: "Say something interesting."}) var chunks []llm.StreamChunk for chunk, err := range stream { diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 6932298..de38251 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -77,8 +77,9 @@ func New(cfg Config) *Loop { // Run executes the agent loop and returns the final response. // Optional history messages are appended after the system prompt and before -// the new user input. -func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (Response, error) { +// the new user input. input's Role is overwritten to RoleUser regardless of +// what the caller sets, so callers only need to fill in Content/Parts. +func (l *Loop) Run(ctx context.Context, input llm.Message, history ...llm.Message) (Response, error) { start := time.Now() messages := l.buildInitialMessages(input, history) @@ -174,8 +175,9 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R // RunStream executes the agent loop with streaming output. // Optional history messages are appended after the system prompt and before -// the new user input. -func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] { +// the new user input. input's Role is overwritten to RoleUser regardless of +// what the caller sets, so callers only need to fill in Content/Parts. +func (l *Loop) RunStream(ctx context.Context, input llm.Message, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] { return func(yield func(llm.StreamChunk, error) bool) { messages := l.buildInitialMessages(input, history) // Same as Run: the schemas are identical on every iteration. @@ -362,12 +364,13 @@ func containsUnparsedToolCall(s string) bool { return strings.Contains(s, "