diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index e64a7afa2..2c2df7dec 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -348,13 +348,14 @@ func TestRun(t *testing.T) { City string `json:"city"` } - weatherTool := agent.FunctionTool[Params]( - "get_weather", - "Get weather for a city", - func(_ context.Context, p Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "Sunny, 22°C in " + p.City}, nil - }, - ) + weatherTool, err := agent.FunctionTool[Params]( + "get_weather", + "Get weather for a city", + func(_ context.Context, p Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "Sunny, 22°C in " + p.City}, nil + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -406,30 +407,31 @@ func TestRun(t *testing.T) { }, } - type Params struct{} - noopTool := agent.FunctionTool[Params]( - "noop", - "No-op", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "ok"}, nil - }, + type Params struct{} + noopTool, err := agent.FunctionTool[Params]( + "noop", + "No-op", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "ok"}, nil + }, + ) + require.NoError(t, err) + + ag := agent.New( + "assistant", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(noopTool), + agent.WithMaxTurns(2), ) - ag := agent.New( - "assistant", - newTestClient(provider), - agent.WithModel("test-model"), - agent.WithTools(noopTool), - agent.WithMaxTurns(2), - ) + _, err = ag.Run( + context.Background(), + []llm.Message{userMessage("loop")}, + ) - _, err := ag.Run( - context.Background(), - []llm.Message{userMessage("loop")}, - ) - - require.Error(t, err) - var maxTurnsErr *agent.MaxTurnsExceededError + require.Error(t, err) + var maxTurnsErr *agent.MaxTurnsExceededError require.ErrorAs(t, err, &maxTurnsErr) assert.Equal(t, 2, maxTurnsErr.MaxTurns) }, @@ -440,16 +442,18 @@ func TestRun(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - makeTool := func(name string) agent.Tool { - return agent.FunctionTool[Params]( - name, - "desc", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "ok"}, nil - }, - ) - } + type Params struct{} + makeTool := func(name string) agent.Tool { + tool, err := agent.FunctionTool[Params]( + name, + "desc", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "ok"}, nil + }, + ) + require.NoError(t, err) + return tool + } provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -546,20 +550,22 @@ func TestRun(t *testing.T) { type Params struct{} - tool1 := agent.FunctionTool[Params]( - "first", - "First tool", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "result_1"}, nil - }, - ) - tool2 := agent.FunctionTool[Params]( - "second", - "Second tool", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "result_2"}, nil - }, - ) + tool1, err := agent.FunctionTool[Params]( + "first", + "First tool", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "result_1"}, nil + }, + ) + require.NoError(t, err) + tool2, err := agent.FunctionTool[Params]( + "second", + "Second tool", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "result_2"}, nil + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -613,20 +619,22 @@ func TestRun(t *testing.T) { type Params struct{} - successTool := agent.FunctionTool[Params]( - "succeed", - "Always succeeds", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "success_result"}, nil - }, - ) - failTool := agent.FunctionTool[Params]( - "fail", - "Always fails", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{}, errors.New("tool exploded") - }, - ) + successTool, err := agent.FunctionTool[Params]( + "succeed", + "Always succeeds", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "success_result"}, nil + }, + ) + require.NoError(t, err) + failTool, err := agent.FunctionTool[Params]( + "fail", + "Always fails", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{}, errors.New("tool exploded") + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -684,16 +692,17 @@ func TestRun(t *testing.T) { var capturedTenantID string - type Params struct{} - tool := agent.FunctionTool[Params]( - "check_tenant", - "Check current tenant", - func(ctx context.Context, _ Params) (agent.ToolResult, error) { - rc := agent.RunContextFrom[*RequestContext](ctx) - capturedTenantID = rc.TenantID - return agent.ToolResult{Content: "tenant: " + rc.TenantID}, nil - }, - ) + type Params struct{} + tool, err := agent.FunctionTool[Params]( + "check_tenant", + "Check current tenant", + func(ctx context.Context, _ Params) (agent.ToolResult, error) { + rc := agent.RunContextFrom[*RequestContext](ctx) + capturedTenantID = rc.TenantID + return agent.ToolResult{Content: "tenant: " + rc.TenantID}, nil + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -1056,26 +1065,27 @@ func TestRun_Hooks(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - noopTool := agent.FunctionTool[Params]( - "noop", - "No-op", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "ok"}, nil - }, - ) + type Params struct{} + noopTool, err := agent.FunctionTool[Params]( + "noop", + "No-op", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "ok"}, nil + }, + ) + require.NoError(t, err) - provider := &mockProvider{ - responses: []*llm.ChatCompletionResponse{ - toolCallResponse(llm.ToolCall{ - ID: "tc_1", - Function: llm.FunctionCall{Name: "noop", Arguments: `{}`}, - }), - stopResponse("done"), - }, - } + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_1", + Function: llm.FunctionCall{Name: "noop", Arguments: `{}`}, + }), + stopResponse("done"), + }, + } - hook := &recordingHook{} + hook := &recordingHook{} ag := agent.New( "assistant", @@ -1340,30 +1350,31 @@ func TestRun_ToolUseBehavior(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - tool := agent.FunctionTool[Params]( - "compute", - "Compute something", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "computed_result"}, nil - }, - ) + type Params struct{} + tool, err := agent.FunctionTool[Params]( + "compute", + "Compute something", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "computed_result"}, nil + }, + ) + require.NoError(t, err) - provider := &mockProvider{ - responses: []*llm.ChatCompletionResponse{ - toolCallResponse(llm.ToolCall{ - ID: "tc_1", - Function: llm.FunctionCall{Name: "compute", Arguments: `{}`}, - }), - }, - } + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_1", + Function: llm.FunctionCall{Name: "compute", Arguments: `{}`}, + }), + }, + } - ag := agent.New( - "assistant", - newTestClient(provider), - agent.WithModel("test-model"), - agent.WithTools(tool), - agent.WithToolUseBehavior(agent.StopOnFirstTool()), + ag := agent.New( + "assistant", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(tool), + agent.WithToolUseBehavior(agent.StopOnFirstTool()), ) result, err := ag.Run( @@ -1385,20 +1396,22 @@ func TestRun_ToolUseBehavior(t *testing.T) { type Params struct{} - tool1 := agent.FunctionTool[Params]( - "search", - "Search", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "search_result"}, nil - }, - ) - tool2 := agent.FunctionTool[Params]( - "submit", - "Submit", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "submitted"}, nil - }, - ) + tool1, err := agent.FunctionTool[Params]( + "search", + "Search", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "search_result"}, nil + }, + ) + require.NoError(t, err) + tool2, err := agent.FunctionTool[Params]( + "submit", + "Submit", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "submitted"}, nil + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -1432,22 +1445,23 @@ func TestRun_ToolUseBehavior(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - tool := agent.FunctionTool[Params]( - "noop", - "No-op", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "ok"}, nil - }, - ) + type Params struct{} + tool, err := agent.FunctionTool[Params]( + "noop", + "No-op", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "ok"}, nil + }, + ) + require.NoError(t, err) - provider := &mockProvider{ - responses: []*llm.ChatCompletionResponse{ - toolCallResponse(llm.ToolCall{ - ID: "tc_1", - Function: llm.FunctionCall{Name: "noop", Arguments: `{}`}, - }), - stopResponse("Final answer."), + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_1", + Function: llm.FunctionCall{Name: "noop", Arguments: `{}`}, + }), + stopResponse("Final answer."), }, } @@ -1474,41 +1488,42 @@ func TestRun_ToolUseBehavior(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - tool := agent.FunctionTool[Params]( - "compute", - "Compute something", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "result"}, nil - }, - ) + type Params struct{} + tool, err := agent.FunctionTool[Params]( + "compute", + "Compute something", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "result"}, nil + }, + ) + require.NoError(t, err) - provider := &mockProvider{ - responses: []*llm.ChatCompletionResponse{ - toolCallResponse(llm.ToolCall{ - ID: "tc_1", - Function: llm.FunctionCall{Name: "compute", Arguments: `{}`}, - }), - }, - } + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_1", + Function: llm.FunctionCall{Name: "compute", Arguments: `{}`}, + }), + }, + } - ag := agent.New( - "assistant", - newTestClient(provider), - agent.WithModel("test-model"), - agent.WithTools(tool), - agent.WithToolUseBehavior(agent.ToolUseBehavior(func(_ context.Context, _ []agent.ToolCallResult) (string, bool, error) { + ag := agent.New( + "assistant", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(tool), + agent.WithToolUseBehavior(agent.ToolUseBehavior(func(_ context.Context, _ []agent.ToolCallResult) (string, bool, error) { return "", false, errors.New("custom behavior failed") })), ) - _, err := ag.Run( - context.Background(), - []llm.Message{userMessage("compute")}, - ) + _, err = ag.Run( + context.Background(), + []llm.Message{userMessage("compute")}, + ) - require.Error(t, err) - assert.Contains(t, err.Error(), "custom behavior failed") + require.Error(t, err) + assert.Contains(t, err.Error(), "custom behavior failed") }, ) } @@ -1527,11 +1542,14 @@ func TestRun_OutputType(t *testing.T) { }, } + infoType, err := agent.NewOutputType[Info]("info") + require.NoError(t, err) + ag := agent.New( "assistant", newTestClient(provider), agent.WithModel("test-model"), - agent.WithOutputType(agent.NewOutputType[Info]("info")), + agent.WithOutputType(infoType), ) result, err := ag.Run( @@ -1560,37 +1578,38 @@ func TestRun_Approval(t *testing.T) { }, } - deleteTool := agent.FunctionTool[struct{}]( - "delete_account", - "Deletes the user account", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - return agent.ToolResult{Content: "deleted"}, nil - }, - ) + deleteTool, err := agent.FunctionTool[struct{}]( + "delete_account", + "Deletes the user account", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + return agent.ToolResult{Content: "deleted"}, nil + }, + ) + require.NoError(t, err) - ag := agent.New( - "assistant", - newTestClient(provider), - agent.WithModel("test-model"), - agent.WithTools(deleteTool), - agent.WithApproval(agent.ApprovalConfig{ - ToolNames: []string{"delete_account"}, - }), - ) + ag := agent.New( + "assistant", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(deleteTool), + agent.WithApproval(agent.ApprovalConfig{ + ToolNames: []string{"delete_account"}, + }), + ) - _, err := ag.Run( - context.Background(), - []llm.Message{userMessage("Delete my account")}, - ) + _, err = ag.Run( + context.Background(), + []llm.Message{userMessage("Delete my account")}, + ) - require.Error(t, err) - var interrupted *agent.InterruptedError - require.ErrorAs(t, err, &interrupted) - assert.Len(t, interrupted.ToolCalls, 1) - assert.Equal(t, "delete_account", interrupted.ToolCalls[0].Function.Name) - assert.Len(t, interrupted.PendingApprovals, 1) - assert.Equal(t, "delete_account", interrupted.PendingApprovals[0].Function.Name) - assert.Equal(t, 1, interrupted.Turns) + require.Error(t, err) + var interrupted *agent.InterruptedError + require.ErrorAs(t, err, &interrupted) + assert.Len(t, interrupted.ToolCalls, 1) + assert.Equal(t, "delete_account", interrupted.ToolCalls[0].Function.Name) + assert.Len(t, interrupted.PendingApprovals, 1) + assert.Equal(t, "delete_account", interrupted.PendingApprovals[0].Function.Name) + assert.Equal(t, 1, interrupted.Turns) }, ) @@ -1599,18 +1618,19 @@ func TestRun_Approval(t *testing.T) { func(t *testing.T) { t.Parallel() - var toolExecuted bool + var toolExecuted bool - deleteTool := agent.FunctionTool[struct{}]( - "delete_account", - "Deletes the user account", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - toolExecuted = true - return agent.ToolResult{Content: "account deleted"}, nil - }, - ) + deleteTool, err := agent.FunctionTool[struct{}]( + "delete_account", + "Deletes the user account", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + toolExecuted = true + return agent.ToolResult{Content: "account deleted"}, nil + }, + ) + require.NoError(t, err) - provider := &mockProvider{ + provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ toolCallResponse(llm.ToolCall{ ID: "tc1", @@ -1630,14 +1650,14 @@ func TestRun_Approval(t *testing.T) { }), ) - _, err := ag.Run( - context.Background(), - []llm.Message{userMessage("Delete my account")}, - ) + _, err = ag.Run( + context.Background(), + []llm.Message{userMessage("Delete my account")}, + ) - var interrupted *agent.InterruptedError - require.ErrorAs(t, err, &interrupted) - assert.False(t, toolExecuted) + var interrupted *agent.InterruptedError + require.ErrorAs(t, err, &interrupted) + assert.False(t, toolExecuted) result, err := agent.Resume( context.Background(), @@ -1660,22 +1680,23 @@ func TestRun_Approval(t *testing.T) { func(t *testing.T) { t.Parallel() - deleteTool := agent.FunctionTool[struct{}]( - "delete_account", - "Deletes the user account", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - t.Fatal("tool should not be executed") - return agent.ToolResult{}, nil - }, - ) + deleteTool, err := agent.FunctionTool[struct{}]( + "delete_account", + "Deletes the user account", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + t.Fatal("tool should not be executed") + return agent.ToolResult{}, nil + }, + ) + require.NoError(t, err) - provider := &mockProvider{ - responses: []*llm.ChatCompletionResponse{ - toolCallResponse(llm.ToolCall{ - ID: "tc1", - Function: llm.FunctionCall{Name: "delete_account", Arguments: `{}`}, - }), - stopResponse("OK, I won't delete your account."), + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc1", + Function: llm.FunctionCall{Name: "delete_account", Arguments: `{}`}, + }), + stopResponse("OK, I won't delete your account."), }, } @@ -1689,7 +1710,7 @@ func TestRun_Approval(t *testing.T) { }), ) - _, err := ag.Run( + _, err = ag.Run( context.Background(), []llm.Message{userMessage("Delete my account")}, ) @@ -1727,21 +1748,22 @@ func TestRun_Approval(t *testing.T) { }, } - safeTool := agent.FunctionTool[struct{}]( - "safe_tool", - "A safe tool", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - return agent.ToolResult{Content: "safe result"}, nil - }, - ) + safeTool, err := agent.FunctionTool[struct{}]( + "safe_tool", + "A safe tool", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + return agent.ToolResult{Content: "safe result"}, nil + }, + ) + require.NoError(t, err) - ag := agent.New( - "assistant", - newTestClient(provider), - agent.WithModel("test-model"), - agent.WithTools(safeTool), - agent.WithApproval(agent.ApprovalConfig{ - ShouldApprove: func(_ context.Context, tc llm.ToolCall) bool { + ag := agent.New( + "assistant", + newTestClient(provider), + agent.WithModel("test-model"), + agent.WithTools(safeTool), + agent.WithApproval(agent.ApprovalConfig{ + ShouldApprove: func(_ context.Context, tc llm.ToolCall) bool { return tc.Function.Name == "dangerous_tool" }, }), @@ -1764,23 +1786,25 @@ func TestRun_Approval(t *testing.T) { var safeExecuted, dangerExecuted bool - safeTool := agent.FunctionTool[struct{}]( - "safe_action", - "A safe action", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - safeExecuted = true - return agent.ToolResult{Content: "safe done"}, nil - }, - ) + safeTool, err := agent.FunctionTool[struct{}]( + "safe_action", + "A safe action", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + safeExecuted = true + return agent.ToolResult{Content: "safe done"}, nil + }, + ) + require.NoError(t, err) - dangerTool := agent.FunctionTool[struct{}]( - "danger_action", - "A dangerous action", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - dangerExecuted = true - return agent.ToolResult{Content: "danger done"}, nil - }, - ) + dangerTool, err := agent.FunctionTool[struct{}]( + "danger_action", + "A dangerous action", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + dangerExecuted = true + return agent.ToolResult{Content: "danger done"}, nil + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -1808,7 +1832,7 @@ func TestRun_Approval(t *testing.T) { }), ) - _, err := ag.Run( + _, err = ag.Run( context.Background(), []llm.Message{userMessage("Do both")}, ) @@ -1916,21 +1940,22 @@ func TestResume(t *testing.T) { func(t *testing.T) { t.Parallel() - deleteTool := agent.FunctionTool[struct{}]( - "delete_account", - "Deletes the user account", - func(_ context.Context, _ struct{}) (agent.ToolResult, error) { - return agent.ToolResult{Content: "deleted"}, nil - }, - ) + deleteTool, err := agent.FunctionTool[struct{}]( + "delete_account", + "Deletes the user account", + func(_ context.Context, _ struct{}) (agent.ToolResult, error) { + return agent.ToolResult{Content: "deleted"}, nil + }, + ) + require.NoError(t, err) - provider := &mockProvider{ - responses: []*llm.ChatCompletionResponse{ - toolCallResponse(llm.ToolCall{ - ID: "tc1", - Function: llm.FunctionCall{Name: "delete_account", Arguments: `{}`}, - }), - stopResponse("Account deleted."), + provider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc1", + Function: llm.FunctionCall{Name: "delete_account", Arguments: `{}`}, + }), + stopResponse("Account deleted."), }, } @@ -1944,7 +1969,7 @@ func TestResume(t *testing.T) { }), ) - _, err := ag.Run( + _, err = ag.Run( context.Background(), []llm.Message{userMessage("Delete my account")}, ) @@ -2123,16 +2148,17 @@ func TestRunStreamed(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - tool := agent.FunctionTool[Params]( - "noop", - "No-op", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "ok"}, nil - }, - ) + type Params struct{} + tool, err := agent.FunctionTool[Params]( + "noop", + "No-op", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "ok"}, nil + }, + ) + require.NoError(t, err) - stream1 := &mockChatStream{ + stream1 := &mockChatStream{ events: []llm.ChatCompletionStreamEvent{ { Delta: llm.MessageDelta{ @@ -2280,21 +2306,23 @@ func TestClone(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - tool1 := agent.FunctionTool[Params]( - "t1", - "desc", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "ok"}, nil - }, - ) - tool2 := agent.FunctionTool[Params]( - "t2", - "desc", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "ok"}, nil - }, - ) + type Params struct{} + tool1, err := agent.FunctionTool[Params]( + "t1", + "desc", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "ok"}, nil + }, + ) + require.NoError(t, err) + tool2, err := agent.FunctionTool[Params]( + "t2", + "desc", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "ok"}, nil + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -2401,13 +2429,14 @@ func TestGenerateSchema_EmbeddedStruct(t *testing.T) { Name string `json:"name"` } - tool := agent.FunctionTool[Params]( + tool, err := agent.FunctionTool[Params]( "create", "Create item", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "ok"}, nil }, ) + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema)) @@ -2575,13 +2604,14 @@ func TestRun_UnknownToolCall(t *testing.T) { } type Params struct{} - tool := agent.FunctionTool[Params]( + tool, err := agent.FunctionTool[Params]( "real_tool", "A real tool", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "ok"}, nil }, ) + require.NoError(t, err) ag := agent.New( "assistant", @@ -2590,7 +2620,7 @@ func TestRun_UnknownToolCall(t *testing.T) { agent.WithTools(tool), ) - _, err := ag.Run( + _, err = ag.Run( context.Background(), []llm.Message{userMessage("test")}, ) @@ -2642,13 +2672,14 @@ func TestClone_WithApprovalConfig(t *testing.T) { }, } - deleteTool := agent.FunctionTool[struct{}]( + deleteTool, err := agent.FunctionTool[struct{}]( "delete", "Delete something", func(_ context.Context, _ struct{}) (agent.ToolResult, error) { return agent.ToolResult{Content: "deleted"}, nil }, ) + require.NoError(t, err) original := agent.New( "assistant", @@ -2662,7 +2693,7 @@ func TestClone_WithApprovalConfig(t *testing.T) { cloned := original.Clone() - _, err := cloned.Run( + _, err = cloned.Run( context.Background(), []llm.Message{userMessage("delete it")}, ) @@ -2692,15 +2723,16 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) { var executionOrder []string - type Params struct{} - tool1 := agent.FunctionTool[Params]( - "prepare", - "Prepare data", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - executionOrder = append(executionOrder, "prepare") - return agent.ToolResult{Content: "prepared"}, nil - }, - ) + type Params struct{} + tool1, err := agent.FunctionTool[Params]( + "prepare", + "Prepare data", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + executionOrder = append(executionOrder, "prepare") + return agent.ToolResult{Content: "prepared"}, nil + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -2751,22 +2783,24 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - tool1 := agent.FunctionTool[Params]( - "prepare", - "Prepare data", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{Content: "prepared"}, nil - }, - ) - tool2 := agent.FunctionTool[Params]( - "finalize", - "Finalize data", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - t.Fatal("tool after handoff should not be executed") - return agent.ToolResult{}, nil - }, - ) + type Params struct{} + tool1, err := agent.FunctionTool[Params]( + "prepare", + "Prepare data", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{Content: "prepared"}, nil + }, + ) + require.NoError(t, err) + tool2, err := agent.FunctionTool[Params]( + "finalize", + "Finalize data", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + t.Fatal("tool after handoff should not be executed") + return agent.ToolResult{}, nil + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -2835,14 +2869,15 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) { func(t *testing.T) { t.Parallel() - type Params struct{} - failingTool := agent.FunctionTool[Params]( - "prepare", - "Prepare data", - func(_ context.Context, _ Params) (agent.ToolResult, error) { - return agent.ToolResult{}, errors.New("preparation failed") - }, - ) + type Params struct{} + failingTool, err := agent.FunctionTool[Params]( + "prepare", + "Prepare data", + func(_ context.Context, _ Params) (agent.ToolResult, error) { + return agent.ToolResult{}, errors.New("preparation failed") + }, + ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -2875,7 +2910,7 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) { agent.WithHandoffs(specialist), ) - _, err := router.Run( + _, err = router.Run( context.Background(), []llm.Message{userMessage("prepare and transfer")}, ) diff --git a/pkg/agent/agent_tool.go b/pkg/agent/agent_tool.go index 0fb112712..16ccad145 100644 --- a/pkg/agent/agent_tool.go +++ b/pkg/agent/agent_tool.go @@ -40,7 +40,7 @@ type ( ) var ( - agentToolParamsSchema = jsonSchemaFor[agentToolParams]() + agentToolParamsSchema = mustJSONSchemaFor[agentToolParams]() ) func agentToolDepth(ctx context.Context) int { diff --git a/pkg/agent/agent_tool_test.go b/pkg/agent/agent_tool_test.go index 4a16e139e..bdcf2447b 100644 --- a/pkg/agent/agent_tool_test.go +++ b/pkg/agent/agent_tool_test.go @@ -252,7 +252,7 @@ func TestAgentTool_Execute(t *testing.T) { var captured string type Params struct{} - tenantTool := agent.FunctionTool[Params]( + tenantTool, err := agent.FunctionTool[Params]( "get_tenant", "Get tenant", func(ctx context.Context, _ Params) (agent.ToolResult, error) { @@ -261,6 +261,7 @@ func TestAgentTool_Execute(t *testing.T) { return agent.ToolResult{Content: rc.TenantID}, nil }, ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -303,13 +304,14 @@ func TestAgentTool_Execute(t *testing.T) { Expr string `json:"expr"` } - calcTool := agent.FunctionTool[Params]( + calcTool, err := agent.FunctionTool[Params]( "calc", "Calculate expression", func(_ context.Context, p Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "42"}, nil }, ) + require.NoError(t, err) provider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -379,13 +381,14 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { func(t *testing.T) { t.Parallel() - deleteTool := agent.FunctionTool[struct{}]( + deleteTool, err := agent.FunctionTool[struct{}]( "delete_file", "Delete a file", func(_ context.Context, _ struct{}) (agent.ToolResult, error) { return agent.ToolResult{Content: "file deleted"}, nil }, ) + require.NoError(t, err) innerProvider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -422,7 +425,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")), ) - _, err := outerAgent.Run( + _, err = outerAgent.Run( context.Background(), []llm.Message{userMessage("Delete the file")}, ) @@ -443,7 +446,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { var toolExecuted bool - deleteTool := agent.FunctionTool[struct{}]( + deleteTool, err := agent.FunctionTool[struct{}]( "delete_file", "Delete a file", func(_ context.Context, _ struct{}) (agent.ToolResult, error) { @@ -451,6 +454,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { return agent.ToolResult{Content: "file deleted"}, nil }, ) + require.NoError(t, err) innerProvider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -489,7 +493,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")), ) - _, err := outerAgent.Run( + _, err = outerAgent.Run( context.Background(), []llm.Message{userMessage("Delete the file")}, ) @@ -520,7 +524,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { func(t *testing.T) { t.Parallel() - deleteTool := agent.FunctionTool[struct{}]( + deleteTool, err := agent.FunctionTool[struct{}]( "delete_file", "Delete a file", func(_ context.Context, _ struct{}) (agent.ToolResult, error) { @@ -528,6 +532,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) innerProvider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -566,7 +571,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")), ) - _, err := outerAgent.Run( + _, err = outerAgent.Run( context.Background(), []llm.Message{userMessage("Delete the file")}, ) @@ -598,7 +603,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { var siblingCalled bool type Params struct{} - siblingTool := agent.FunctionTool[Params]( + siblingTool, err := agent.FunctionTool[Params]( "list_files", "List files", func(_ context.Context, _ Params) (agent.ToolResult, error) { @@ -606,14 +611,16 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { return agent.ToolResult{Content: "file1.txt, file2.txt"}, nil }, ) + require.NoError(t, err) - deleteTool := agent.FunctionTool[struct{}]( + deleteTool, err := agent.FunctionTool[struct{}]( "delete_file", "Delete a file", func(_ context.Context, _ struct{}) (agent.ToolResult, error) { return agent.ToolResult{Content: "file deleted"}, nil }, ) + require.NoError(t, err) innerProvider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -661,7 +668,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { ), ) - _, err := outerAgent.Run( + _, err = outerAgent.Run( context.Background(), []llm.Message{userMessage("List and delete files")}, ) @@ -692,7 +699,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { var toolExecuted bool - dangerTool := agent.FunctionTool[struct{}]( + dangerTool, err := agent.FunctionTool[struct{}]( "danger", "Dangerous operation", func(_ context.Context, _ struct{}) (agent.ToolResult, error) { @@ -700,6 +707,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { return agent.ToolResult{Content: "danger executed"}, nil }, ) + require.NoError(t, err) cProvider := &mockProvider{ responses: []*llm.ChatCompletionResponse{ @@ -755,7 +763,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { agent.WithTools(agentB.AsTool("call_b", "Call agent B")), ) - _, err := agentA.Run( + _, err = agentA.Run( context.Background(), []llm.Message{userMessage("start")}, ) diff --git a/pkg/agent/handoff.go b/pkg/agent/handoff.go index cec8c83eb..a5d348d9d 100644 --- a/pkg/agent/handoff.go +++ b/pkg/agent/handoff.go @@ -45,7 +45,7 @@ type ( ) var ( - handoffParamsSchema = jsonSchemaFor[handoffParams]() + handoffParamsSchema = mustJSONSchemaFor[handoffParams]() ) func HandoffTo(agent *Agent, opts ...HandoffOption) *Handoff { diff --git a/pkg/agent/output_type.go b/pkg/agent/output_type.go index 160495682..d99c132ed 100644 --- a/pkg/agent/output_type.go +++ b/pkg/agent/output_type.go @@ -16,6 +16,7 @@ package agent import ( "encoding/json" + "fmt" "go.probo.inc/probo/pkg/llm" ) @@ -27,11 +28,16 @@ type OutputType struct { Schema json.RawMessage } -func NewOutputType[T any](name string) *OutputType { +func NewOutputType[T any](name string) (*OutputType, error) { + schema, err := jsonSchemaFor[T]() + if err != nil { + return nil, fmt.Errorf("cannot create output type %q: %w", name, err) + } + return &OutputType{ Name: name, - Schema: jsonSchemaFor[T](), - } + Schema: schema, + }, nil } func (o *OutputType) responseFormat() *llm.ResponseFormat { diff --git a/pkg/agent/output_type_test.go b/pkg/agent/output_type_test.go index d411e96ac..d2846b9ca 100644 --- a/pkg/agent/output_type_test.go +++ b/pkg/agent/output_type_test.go @@ -31,7 +31,8 @@ func TestNewOutputType_SetsNameAndSchema(t *testing.T) { Score int `json:"score"` } - ot := NewOutputType[Result]("test_result") + ot, err := NewOutputType[Result]("test_result") + require.NoError(t, err) assert.Equal(t, "test_result", ot.Name) require.NotNil(t, ot.Schema) @@ -50,7 +51,8 @@ func TestNewOutputType_EmptyStruct(t *testing.T) { type Empty struct{} - ot := NewOutputType[Empty]("empty") + ot, err := NewOutputType[Empty]("empty") + require.NoError(t, err) assert.Equal(t, "empty", ot.Name) @@ -71,7 +73,8 @@ func TestOutputType_responseFormat(t *testing.T) { Reason string `json:"reason"` } - ot := NewOutputType[Verdict]("verdict") + ot, err := NewOutputType[Verdict]("verdict") + require.NoError(t, err) rf := ot.responseFormat() require.NotNil(t, rf) @@ -91,7 +94,8 @@ func TestOutputType_responseFormat_SchemaMatchesOutputType(t *testing.T) { Priority *int `json:"priority,omitempty"` } - ot := NewOutputType[Analysis]("analysis") + ot, err := NewOutputType[Analysis]("analysis") + require.NoError(t, err) rf := ot.responseFormat() var schema map[string]any diff --git a/pkg/agent/schema.go b/pkg/agent/schema.go index 93eabf43c..12ad5e930 100644 --- a/pkg/agent/schema.go +++ b/pkg/agent/schema.go @@ -22,22 +22,30 @@ import ( "github.com/google/jsonschema-go/jsonschema" ) -func jsonSchemaFor[T any]() json.RawMessage { +func jsonSchemaFor[T any]() (json.RawMessage, error) { t := reflect.TypeFor[T]() schema, err := jsonschema.ForType(t, nil) if err != nil { - panic(fmt.Sprintf("cannot generate schema for %s: %v", t, err)) + return nil, fmt.Errorf("cannot generate schema for %s: %w", t, err) } stripNullTypes(schema) data, err := json.Marshal(schema) if err != nil { - panic(fmt.Sprintf("cannot marshal schema for %s: %v", t, err)) + return nil, fmt.Errorf("cannot marshal schema for %s: %w", t, err) } - return json.RawMessage(data) + return json.RawMessage(data), nil +} + +func mustJSONSchemaFor[T any]() json.RawMessage { + schema, err := jsonSchemaFor[T]() + if err != nil { + panic(err) + } + return schema } // stripNullTypes removes "null" from union types produced by pointer fields diff --git a/pkg/agent/schema_test.go b/pkg/agent/schema_test.go index f0b95f832..6796c865d 100644 --- a/pkg/agent/schema_test.go +++ b/pkg/agent/schema_test.go @@ -32,7 +32,8 @@ func TestGenerateSchema_PointerFieldsStripNull(t *testing.T) { Done *bool `json:"done"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -65,7 +66,8 @@ func TestGenerateSchema_IntegerBoundsStripped(t *testing.T) { Uint16 uint16 `json:"uint16"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -85,7 +87,8 @@ func TestGenerateSchema_EmptyStruct(t *testing.T) { type Params struct{} - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -104,7 +107,8 @@ func TestGenerateSchema_MapField(t *testing.T) { Metadata map[string]string `json:"metadata"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -129,7 +133,8 @@ func TestGenerateSchema_NestedPointerStruct(t *testing.T) { Inner *Inner `json:"inner"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -153,7 +158,8 @@ func TestGenerateSchema_SliceOfPointers(t *testing.T) { Names []*string `json:"names"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -175,7 +181,8 @@ func TestGenerateSchema_MapWithPointerValues(t *testing.T) { Scores map[string]*int `json:"scores"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -200,7 +207,8 @@ func TestGenerateSchema_DescriptionFromJsonschemaTag(t *testing.T) { Limit int `json:"limit" jsonschema:"Maximum number of results to return"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -224,7 +232,8 @@ func TestGenerateSchema_RequiredVsOptional(t *testing.T) { AlsoNeeded int `json:"also_needed"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -252,7 +261,8 @@ func TestGenerateSchema_DeeplyNestedStructure(t *testing.T) { Root Level1 `json:"root"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) @@ -293,7 +303,8 @@ func TestGenerateSchema_SliceOfStructs(t *testing.T) { Items []Item `json:"items"` } - raw := jsonSchemaFor[Params]() + raw, err := jsonSchemaFor[Params]() + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(raw, &schema)) diff --git a/pkg/agent/tool.go b/pkg/agent/tool.go index 09cf0eb8c..daf0613b2 100644 --- a/pkg/agent/tool.go +++ b/pkg/agent/tool.go @@ -52,8 +52,11 @@ func FunctionTool[P any]( name string, description string, fn func(ctx context.Context, params P) (ToolResult, error), -) Tool { - schema := jsonSchemaFor[P]() +) (Tool, error) { + schema, err := jsonSchemaFor[P]() + if err != nil { + return nil, fmt.Errorf("cannot create tool %q: %w", name, err) + } var parsed struct { Required []string `json:"required"` @@ -66,7 +69,7 @@ func FunctionTool[P any]( fn: fn, schema: schema, requiredFields: parsed.Required, - } + }, nil } func (t *functionTool[P]) Name() string { return t.name } diff --git a/pkg/agent/tool_test.go b/pkg/agent/tool_test.go index 76cee47eb..4a9b8a4d8 100644 --- a/pkg/agent/tool_test.go +++ b/pkg/agent/tool_test.go @@ -30,13 +30,14 @@ func TestFunctionTool_Name(t *testing.T) { type Params struct{} - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "my_tool", "does things", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) assert.Equal(t, "my_tool", tool.Name()) } @@ -53,13 +54,14 @@ func TestFunctionTool_Definition(t *testing.T) { Query string `json:"query"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "search", "Search for items", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) def := tool.Definition() assert.Equal(t, "search", def.Name) @@ -77,13 +79,14 @@ func TestFunctionTool_Definition(t *testing.T) { Count int `json:"count"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "create", "Create items", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) def := tool.Definition() require.NotNil(t, def.Parameters) @@ -113,13 +116,14 @@ func TestFunctionTool_Definition(t *testing.T) { type Params struct{} - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "noop", "No-op", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema)) @@ -136,13 +140,14 @@ func TestFunctionTool_Definition(t *testing.T) { Title *string `json:"title,omitempty"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "update", "Update", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) var schema map[string]any require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema)) @@ -168,13 +173,14 @@ func TestFunctionTool_Execute(t *testing.T) { Y int `json:"y"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "add", "Add two numbers", func(_ context.Context, p Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "42"}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{"x": 1, "y": 2}`) require.NoError(t, err) @@ -193,7 +199,7 @@ func TestFunctionTool_Execute(t *testing.T) { } var received string - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "weather", "Get weather", func(_ context.Context, p Params) (agent.ToolResult, error) { @@ -201,8 +207,9 @@ func TestFunctionTool_Execute(t *testing.T) { return agent.ToolResult{Content: "sunny"}, nil }, ) + require.NoError(t, err) - _, err := tool.Execute(context.Background(), `{"city":"Paris"}`) + _, err = tool.Execute(context.Background(), `{"city":"Paris"}`) require.NoError(t, err) assert.Equal(t, "Paris", received) }, @@ -215,13 +222,14 @@ func TestFunctionTool_Execute(t *testing.T) { type Params struct{} - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "noop", "No-op", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "ok"}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{invalid`) require.NoError(t, err) @@ -237,15 +245,16 @@ func TestFunctionTool_Execute(t *testing.T) { type Params struct{} - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "fail", "Always fails", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{}, errors.New("db down") }, ) + require.NoError(t, err) - _, err := tool.Execute(context.Background(), `{}`) + _, err = tool.Execute(context.Background(), `{}`) require.Error(t, err) assert.Contains(t, err.Error(), "db down") }, @@ -259,7 +268,7 @@ func TestFunctionTool_Execute(t *testing.T) { type ctxKey struct{} type Params struct{} - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "ctx_check", "Check context", func(ctx context.Context, _ Params) (agent.ToolResult, error) { @@ -267,6 +276,7 @@ func TestFunctionTool_Execute(t *testing.T) { return agent.ToolResult{Content: val}, nil }, ) + require.NoError(t, err) ctx := context.WithValue(context.Background(), ctxKey{}, "hello") result, err := tool.Execute(ctx, `{}`) @@ -284,7 +294,7 @@ func TestFunctionTool_Execute(t *testing.T) { City string `json:"city"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "weather", "Get weather", func(_ context.Context, _ Params) (agent.ToolResult, error) { @@ -292,6 +302,7 @@ func TestFunctionTool_Execute(t *testing.T) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{}`) require.NoError(t, err) @@ -310,7 +321,7 @@ func TestFunctionTool_Execute(t *testing.T) { Country string `json:"country"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "weather", "Get weather", func(_ context.Context, _ Params) (agent.ToolResult, error) { @@ -318,6 +329,7 @@ func TestFunctionTool_Execute(t *testing.T) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{}`) require.NoError(t, err) @@ -337,7 +349,7 @@ func TestFunctionTool_Execute(t *testing.T) { Country string `json:"country"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "weather", "Get weather", func(_ context.Context, _ Params) (agent.ToolResult, error) { @@ -345,6 +357,7 @@ func TestFunctionTool_Execute(t *testing.T) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{"city":"Paris"}`) require.NoError(t, err) @@ -363,13 +376,14 @@ func TestFunctionTool_Execute(t *testing.T) { Units *string `json:"units,omitempty"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "weather", "Get weather", func(_ context.Context, p Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "sunny in " + p.City}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{"city":"Paris"}`) require.NoError(t, err) @@ -387,13 +401,14 @@ func TestFunctionTool_Execute(t *testing.T) { Name string `json:"name"` } - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "greet", "Greet", func(_ context.Context, p Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "hi " + p.Name}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{"name":"Alice","extra":"ignored"}`) require.NoError(t, err) @@ -409,13 +424,14 @@ func TestFunctionTool_Execute(t *testing.T) { type Params struct{} - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "ping", "Ping", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "pong"}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{}`) require.NoError(t, err) @@ -430,13 +446,14 @@ func TestFunctionTool_Execute(t *testing.T) { type Params struct{} - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "validate", "Validate input", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{Content: "validation failed", IsError: true}, nil }, ) + require.NoError(t, err) result, err := tool.Execute(context.Background(), `{}`) require.NoError(t, err) @@ -451,13 +468,14 @@ func TestFunctionTool_InterfaceSatisfaction(t *testing.T) { type Params struct{} - tool := agent.FunctionTool( + tool, err := agent.FunctionTool( "test", "test tool", func(_ context.Context, _ Params) (agent.ToolResult, error) { return agent.ToolResult{}, nil }, ) + require.NoError(t, err) assert.Implements(t, (*agent.Tool)(nil), tool) assert.Implements(t, (*agent.ToolDescriptor)(nil), tool) diff --git a/pkg/agent/typed.go b/pkg/agent/typed.go index fd3a8fe04..ee8e81b9c 100644 --- a/pkg/agent/typed.go +++ b/pkg/agent/typed.go @@ -36,11 +36,16 @@ func RunTyped[T any]( ) (*TypedResult[T], error) { typed := a.clone() + schema, err := jsonSchemaFor[T]() + if err != nil { + return nil, fmt.Errorf("cannot create typed runner: %w", err) + } + typed.responseFormat = &llm.ResponseFormat{ Type: llm.ResponseFormatJSONSchema, JSONSchema: &llm.JSONSchema{ Name: typeName[T](), - Schema: jsonSchemaFor[T](), + Schema: schema, Strict: true, }, } diff --git a/pkg/agent/typed_test.go b/pkg/agent/typed_test.go index eed59b940..c81f8939e 100644 --- a/pkg/agent/typed_test.go +++ b/pkg/agent/typed_test.go @@ -404,13 +404,14 @@ func TestRunTyped(t *testing.T) { City string `json:"city"` } - weatherTool := FunctionTool[Params]( + weatherTool, err := FunctionTool[Params]( "get_weather", "Get weather for a city", func(_ context.Context, p Params) (ToolResult, error) { return ToolResult{Content: "Sunny, 22°C in " + p.City}, nil }, ) + require.NoError(t, err) provider := &typedMockProvider{ responses: []*llm.ChatCompletionResponse{