rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Return errors from schema generation instead of panicking

Browse Source 
jsonSchemaFor panicked on unsupported types, which meant
FunctionTool, NewOutputType, and RunTyped would crash the
process during setup rather than returning a normal error.

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2026-03-13 18:56:56 +01:00
parent def8f417ca
commit f3239a1a7b
12 changed files with 505 additions and 406 deletions
Download Patch File Download Diff File Expand all files Collapse all files

723
pkg/agent/agent_test.go
View File

@@ -348,13 +348,14 @@ func TestRun(t *testing.T) {
City string `json:"city"` City string `json:"city"`
} }
weatherTool := agent.FunctionTool[Params]( weatherTool, err := agent.FunctionTool[Params](
"get_weather", "get_weather",
"Get weather for a city", "Get weather for a city",
func(_ context.Context, p Params) (agent.ToolResult, error) { func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "Sunny, 22°C in " + p.City}, nil return agent.ToolResult{Content: "Sunny, 22°C in " + p.City}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -406,30 +407,31 @@ func TestRun(t *testing.T) {
}, },
} }
type Params struct{} type Params struct{}
noopTool := agent.FunctionTool[Params]( noopTool, err := agent.FunctionTool[Params](
"noop", "noop",
"No-op", "No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil 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( _, err = ag.Run(
"assistant", context.Background(),
newTestClient(provider), []llm.Message{userMessage("loop")},
agent.WithModel("test-model"), )
agent.WithTools(noopTool),
agent.WithMaxTurns(2),
)
_, err := ag.Run( require.Error(t, err)
context.Background(), var maxTurnsErr *agent.MaxTurnsExceededError
[]llm.Message{userMessage("loop")},
)
require.Error(t, err)
var maxTurnsErr *agent.MaxTurnsExceededError
require.ErrorAs(t, err, &maxTurnsErr) require.ErrorAs(t, err, &maxTurnsErr)
assert.Equal(t, 2, maxTurnsErr.MaxTurns) assert.Equal(t, 2, maxTurnsErr.MaxTurns)
}, },
@@ -440,16 +442,18 @@ func TestRun(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
makeTool := func(name string) agent.Tool { makeTool := func(name string) agent.Tool {
return agent.FunctionTool[Params]( tool, err := agent.FunctionTool[Params](
name, name,
"desc", "desc",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil return agent.ToolResult{Content: "ok"}, nil
}, },
) )
} require.NoError(t, err)
return tool
}
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -546,20 +550,22 @@ func TestRun(t *testing.T) {
type Params struct{} type Params struct{}
tool1 := agent.FunctionTool[Params]( tool1, err := agent.FunctionTool[Params](
"first", "first",
"First tool", "First tool",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "result_1"}, nil return agent.ToolResult{Content: "result_1"}, nil
}, },
) )
tool2 := agent.FunctionTool[Params]( require.NoError(t, err)
"second", tool2, err := agent.FunctionTool[Params](
"Second tool", "second",
func(_ context.Context, _ Params) (agent.ToolResult, error) { "Second tool",
return agent.ToolResult{Content: "result_2"}, nil func(_ context.Context, _ Params) (agent.ToolResult, error) {
}, return agent.ToolResult{Content: "result_2"}, nil
) },
)
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -613,20 +619,22 @@ func TestRun(t *testing.T) {
type Params struct{} type Params struct{}
successTool := agent.FunctionTool[Params]( successTool, err := agent.FunctionTool[Params](
"succeed", "succeed",
"Always succeeds", "Always succeeds",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "success_result"}, nil return agent.ToolResult{Content: "success_result"}, nil
}, },
) )
failTool := agent.FunctionTool[Params]( require.NoError(t, err)
"fail", failTool, err := agent.FunctionTool[Params](
"Always fails", "fail",
func(_ context.Context, _ Params) (agent.ToolResult, error) { "Always fails",
return agent.ToolResult{}, errors.New("tool exploded") func(_ context.Context, _ Params) (agent.ToolResult, error) {
}, return agent.ToolResult{}, errors.New("tool exploded")
) },
)
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -684,16 +692,17 @@ func TestRun(t *testing.T) {
var capturedTenantID string var capturedTenantID string
type Params struct{} type Params struct{}
tool := agent.FunctionTool[Params]( tool, err := agent.FunctionTool[Params](
"check_tenant", "check_tenant",
"Check current tenant", "Check current tenant",
func(ctx context.Context, _ Params) (agent.ToolResult, error) { func(ctx context.Context, _ Params) (agent.ToolResult, error) {
rc := agent.RunContextFrom[*RequestContext](ctx) rc := agent.RunContextFrom[*RequestContext](ctx)
capturedTenantID = rc.TenantID capturedTenantID = rc.TenantID
return agent.ToolResult{Content: "tenant: " + rc.TenantID}, nil return agent.ToolResult{Content: "tenant: " + rc.TenantID}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -1056,26 +1065,27 @@ func TestRun_Hooks(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
noopTool := agent.FunctionTool[Params]( noopTool, err := agent.FunctionTool[Params](
"noop", "noop",
"No-op", "No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil return agent.ToolResult{Content: "ok"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
toolCallResponse(llm.ToolCall{ toolCallResponse(llm.ToolCall{
ID: "tc_1", ID: "tc_1",
Function: llm.FunctionCall{Name: "noop", Arguments: `{}`}, Function: llm.FunctionCall{Name: "noop", Arguments: `{}`},
}), }),
stopResponse("done"), stopResponse("done"),
}, },
} }
hook := &recordingHook{} hook := &recordingHook{}
ag := agent.New( ag := agent.New(
"assistant", "assistant",
@@ -1340,30 +1350,31 @@ func TestRun_ToolUseBehavior(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
tool := agent.FunctionTool[Params]( tool, err := agent.FunctionTool[Params](
"compute", "compute",
"Compute something", "Compute something",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "computed_result"}, nil return agent.ToolResult{Content: "computed_result"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
toolCallResponse(llm.ToolCall{ toolCallResponse(llm.ToolCall{
ID: "tc_1", ID: "tc_1",
Function: llm.FunctionCall{Name: "compute", Arguments: `{}`}, Function: llm.FunctionCall{Name: "compute", Arguments: `{}`},
}), }),
}, },
} }
ag := agent.New( ag := agent.New(
"assistant", "assistant",
newTestClient(provider), newTestClient(provider),
agent.WithModel("test-model"), agent.WithModel("test-model"),
agent.WithTools(tool), agent.WithTools(tool),
agent.WithToolUseBehavior(agent.StopOnFirstTool()), agent.WithToolUseBehavior(agent.StopOnFirstTool()),
) )
result, err := ag.Run( result, err := ag.Run(
@@ -1385,20 +1396,22 @@ func TestRun_ToolUseBehavior(t *testing.T) {
type Params struct{} type Params struct{}
tool1 := agent.FunctionTool[Params]( tool1, err := agent.FunctionTool[Params](
"search", "search",
"Search", "Search",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "search_result"}, nil return agent.ToolResult{Content: "search_result"}, nil
}, },
) )
tool2 := agent.FunctionTool[Params]( require.NoError(t, err)
"submit", tool2, err := agent.FunctionTool[Params](
"Submit", "submit",
func(_ context.Context, _ Params) (agent.ToolResult, error) { "Submit",
return agent.ToolResult{Content: "submitted"}, nil func(_ context.Context, _ Params) (agent.ToolResult, error) {
}, return agent.ToolResult{Content: "submitted"}, nil
) },
)
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -1432,22 +1445,23 @@ func TestRun_ToolUseBehavior(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
tool := agent.FunctionTool[Params]( tool, err := agent.FunctionTool[Params](
"noop", "noop",
"No-op", "No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil return agent.ToolResult{Content: "ok"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
toolCallResponse(llm.ToolCall{ toolCallResponse(llm.ToolCall{
ID: "tc_1", ID: "tc_1",
Function: llm.FunctionCall{Name: "noop", Arguments: `{}`}, Function: llm.FunctionCall{Name: "noop", Arguments: `{}`},
}), }),
stopResponse("Final answer."), stopResponse("Final answer."),
}, },
} }
@@ -1474,41 +1488,42 @@ func TestRun_ToolUseBehavior(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
tool := agent.FunctionTool[Params]( tool, err := agent.FunctionTool[Params](
"compute", "compute",
"Compute something", "Compute something",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "result"}, nil return agent.ToolResult{Content: "result"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
toolCallResponse(llm.ToolCall{ toolCallResponse(llm.ToolCall{
ID: "tc_1", ID: "tc_1",
Function: llm.FunctionCall{Name: "compute", Arguments: `{}`}, Function: llm.FunctionCall{Name: "compute", Arguments: `{}`},
}), }),
}, },
} }
ag := agent.New( ag := agent.New(
"assistant", "assistant",
newTestClient(provider), newTestClient(provider),
agent.WithModel("test-model"), agent.WithModel("test-model"),
agent.WithTools(tool), agent.WithTools(tool),
agent.WithToolUseBehavior(agent.ToolUseBehavior(func(_ context.Context, _ []agent.ToolCallResult) (string, bool, error) { agent.WithToolUseBehavior(agent.ToolUseBehavior(func(_ context.Context, _ []agent.ToolCallResult) (string, bool, error) {
return "", false, errors.New("custom behavior failed") return "", false, errors.New("custom behavior failed")
})), })),
) )
_, err := ag.Run( _, err = ag.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("compute")}, []llm.Message{userMessage("compute")},
) )
require.Error(t, err) require.Error(t, err)
assert.Contains(t, err.Error(), "custom behavior failed") 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( ag := agent.New(
"assistant", "assistant",
newTestClient(provider), newTestClient(provider),
agent.WithModel("test-model"), agent.WithModel("test-model"),
agent.WithOutputType(agent.NewOutputType[Info]("info")), agent.WithOutputType(infoType),
) )
result, err := ag.Run( result, err := ag.Run(
@@ -1560,37 +1578,38 @@ func TestRun_Approval(t *testing.T) {
}, },
} }
deleteTool := agent.FunctionTool[struct{}]( deleteTool, err := agent.FunctionTool[struct{}](
"delete_account", "delete_account",
"Deletes the user account", "Deletes the user account",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "deleted"}, nil return agent.ToolResult{Content: "deleted"}, nil
}, },
) )
require.NoError(t, err)
ag := agent.New( ag := agent.New(
"assistant", "assistant",
newTestClient(provider), newTestClient(provider),
agent.WithModel("test-model"), agent.WithModel("test-model"),
agent.WithTools(deleteTool), agent.WithTools(deleteTool),
agent.WithApproval(agent.ApprovalConfig{ agent.WithApproval(agent.ApprovalConfig{
ToolNames: []string{"delete_account"}, ToolNames: []string{"delete_account"},
}), }),
) )
_, err := ag.Run( _, err = ag.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("Delete my account")}, []llm.Message{userMessage("Delete my account")},
) )
require.Error(t, err) require.Error(t, err)
var interrupted *agent.InterruptedError var interrupted *agent.InterruptedError
require.ErrorAs(t, err, &interrupted) require.ErrorAs(t, err, &interrupted)
assert.Len(t, interrupted.ToolCalls, 1) assert.Len(t, interrupted.ToolCalls, 1)
assert.Equal(t, "delete_account", interrupted.ToolCalls[0].Function.Name) assert.Equal(t, "delete_account", interrupted.ToolCalls[0].Function.Name)
assert.Len(t, interrupted.PendingApprovals, 1) assert.Len(t, interrupted.PendingApprovals, 1)
assert.Equal(t, "delete_account", interrupted.PendingApprovals[0].Function.Name) assert.Equal(t, "delete_account", interrupted.PendingApprovals[0].Function.Name)
assert.Equal(t, 1, interrupted.Turns) assert.Equal(t, 1, interrupted.Turns)
}, },
) )
@@ -1599,18 +1618,19 @@ func TestRun_Approval(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
var toolExecuted bool var toolExecuted bool
deleteTool := agent.FunctionTool[struct{}]( deleteTool, err := agent.FunctionTool[struct{}](
"delete_account", "delete_account",
"Deletes the user account", "Deletes the user account",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
toolExecuted = true toolExecuted = true
return agent.ToolResult{Content: "account deleted"}, nil return agent.ToolResult{Content: "account deleted"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
toolCallResponse(llm.ToolCall{ toolCallResponse(llm.ToolCall{
ID: "tc1", ID: "tc1",
@@ -1630,14 +1650,14 @@ func TestRun_Approval(t *testing.T) {
}), }),
) )
_, err := ag.Run( _, err = ag.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("Delete my account")}, []llm.Message{userMessage("Delete my account")},
) )
var interrupted *agent.InterruptedError var interrupted *agent.InterruptedError
require.ErrorAs(t, err, &interrupted) require.ErrorAs(t, err, &interrupted)
assert.False(t, toolExecuted) assert.False(t, toolExecuted)
result, err := agent.Resume( result, err := agent.Resume(
context.Background(), context.Background(),
@@ -1660,22 +1680,23 @@ func TestRun_Approval(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
deleteTool := agent.FunctionTool[struct{}]( deleteTool, err := agent.FunctionTool[struct{}](
"delete_account", "delete_account",
"Deletes the user account", "Deletes the user account",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
t.Fatal("tool should not be executed") t.Fatal("tool should not be executed")
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
toolCallResponse(llm.ToolCall{ toolCallResponse(llm.ToolCall{
ID: "tc1", ID: "tc1",
Function: llm.FunctionCall{Name: "delete_account", Arguments: `{}`}, Function: llm.FunctionCall{Name: "delete_account", Arguments: `{}`},
}), }),
stopResponse("OK, I won't delete your account."), 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(), context.Background(),
[]llm.Message{userMessage("Delete my account")}, []llm.Message{userMessage("Delete my account")},
) )
@@ -1727,21 +1748,22 @@ func TestRun_Approval(t *testing.T) {
}, },
} }
safeTool := agent.FunctionTool[struct{}]( safeTool, err := agent.FunctionTool[struct{}](
"safe_tool", "safe_tool",
"A safe tool", "A safe tool",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "safe result"}, nil return agent.ToolResult{Content: "safe result"}, nil
}, },
) )
require.NoError(t, err)
ag := agent.New( ag := agent.New(
"assistant", "assistant",
newTestClient(provider), newTestClient(provider),
agent.WithModel("test-model"), agent.WithModel("test-model"),
agent.WithTools(safeTool), agent.WithTools(safeTool),
agent.WithApproval(agent.ApprovalConfig{ agent.WithApproval(agent.ApprovalConfig{
ShouldApprove: func(_ context.Context, tc llm.ToolCall) bool { ShouldApprove: func(_ context.Context, tc llm.ToolCall) bool {
return tc.Function.Name == "dangerous_tool" return tc.Function.Name == "dangerous_tool"
}, },
}), }),
@@ -1764,23 +1786,25 @@ func TestRun_Approval(t *testing.T) {
var safeExecuted, dangerExecuted bool var safeExecuted, dangerExecuted bool
safeTool := agent.FunctionTool[struct{}]( safeTool, err := agent.FunctionTool[struct{}](
"safe_action", "safe_action",
"A safe action", "A safe action",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
safeExecuted = true safeExecuted = true
return agent.ToolResult{Content: "safe done"}, nil return agent.ToolResult{Content: "safe done"}, nil
}, },
) )
require.NoError(t, err)
dangerTool := agent.FunctionTool[struct{}]( dangerTool, err := agent.FunctionTool[struct{}](
"danger_action", "danger_action",
"A dangerous action", "A dangerous action",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
dangerExecuted = true dangerExecuted = true
return agent.ToolResult{Content: "danger done"}, nil return agent.ToolResult{Content: "danger done"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -1808,7 +1832,7 @@ func TestRun_Approval(t *testing.T) {
}), }),
) )
_, err := ag.Run( _, err = ag.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("Do both")}, []llm.Message{userMessage("Do both")},
) )
@@ -1916,21 +1940,22 @@ func TestResume(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
deleteTool := agent.FunctionTool[struct{}]( deleteTool, err := agent.FunctionTool[struct{}](
"delete_account", "delete_account",
"Deletes the user account", "Deletes the user account",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "deleted"}, nil return agent.ToolResult{Content: "deleted"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
toolCallResponse(llm.ToolCall{ toolCallResponse(llm.ToolCall{
ID: "tc1", ID: "tc1",
Function: llm.FunctionCall{Name: "delete_account", Arguments: `{}`}, Function: llm.FunctionCall{Name: "delete_account", Arguments: `{}`},
}), }),
stopResponse("Account deleted."), stopResponse("Account deleted."),
}, },
} }
@@ -1944,7 +1969,7 @@ func TestResume(t *testing.T) {
}), }),
) )
_, err := ag.Run( _, err = ag.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("Delete my account")}, []llm.Message{userMessage("Delete my account")},
) )
@@ -2123,16 +2148,17 @@ func TestRunStreamed(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
tool := agent.FunctionTool[Params]( tool, err := agent.FunctionTool[Params](
"noop", "noop",
"No-op", "No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil return agent.ToolResult{Content: "ok"}, nil
}, },
) )
require.NoError(t, err)
stream1 := &mockChatStream{ stream1 := &mockChatStream{
events: []llm.ChatCompletionStreamEvent{ events: []llm.ChatCompletionStreamEvent{
{ {
Delta: llm.MessageDelta{ Delta: llm.MessageDelta{
@@ -2280,21 +2306,23 @@ func TestClone(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
tool1 := agent.FunctionTool[Params]( tool1, err := agent.FunctionTool[Params](
"t1", "t1",
"desc", "desc",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil return agent.ToolResult{Content: "ok"}, nil
}, },
) )
tool2 := agent.FunctionTool[Params]( require.NoError(t, err)
"t2", tool2, err := agent.FunctionTool[Params](
"desc", "t2",
func(_ context.Context, _ Params) (agent.ToolResult, error) { "desc",
return agent.ToolResult{Content: "ok"}, nil func(_ context.Context, _ Params) (agent.ToolResult, error) {
}, return agent.ToolResult{Content: "ok"}, nil
) },
)
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -2401,13 +2429,14 @@ func TestGenerateSchema_EmbeddedStruct(t *testing.T) {
Name string `json:"name"` Name string `json:"name"`
} }
tool := agent.FunctionTool[Params]( tool, err := agent.FunctionTool[Params](
"create", "create",
"Create item", "Create item",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil return agent.ToolResult{Content: "ok"}, nil
}, },
) )
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema)) require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema))
@@ -2575,13 +2604,14 @@ func TestRun_UnknownToolCall(t *testing.T) {
} }
type Params struct{} type Params struct{}
tool := agent.FunctionTool[Params]( tool, err := agent.FunctionTool[Params](
"real_tool", "real_tool",
"A real tool", "A real tool",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil return agent.ToolResult{Content: "ok"}, nil
}, },
) )
require.NoError(t, err)
ag := agent.New( ag := agent.New(
"assistant", "assistant",
@@ -2590,7 +2620,7 @@ func TestRun_UnknownToolCall(t *testing.T) {
agent.WithTools(tool), agent.WithTools(tool),
) )
_, err := ag.Run( _, err = ag.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("test")}, []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",
"Delete something", "Delete something",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "deleted"}, nil return agent.ToolResult{Content: "deleted"}, nil
}, },
) )
require.NoError(t, err)
original := agent.New( original := agent.New(
"assistant", "assistant",
@@ -2662,7 +2693,7 @@ func TestClone_WithApprovalConfig(t *testing.T) {
cloned := original.Clone() cloned := original.Clone()
_, err := cloned.Run( _, err = cloned.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("delete it")}, []llm.Message{userMessage("delete it")},
) )
@@ -2692,15 +2723,16 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
var executionOrder []string var executionOrder []string
type Params struct{} type Params struct{}
tool1 := agent.FunctionTool[Params]( tool1, err := agent.FunctionTool[Params](
"prepare", "prepare",
"Prepare data", "Prepare data",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
executionOrder = append(executionOrder, "prepare") executionOrder = append(executionOrder, "prepare")
return agent.ToolResult{Content: "prepared"}, nil return agent.ToolResult{Content: "prepared"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -2751,22 +2783,24 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
tool1 := agent.FunctionTool[Params]( tool1, err := agent.FunctionTool[Params](
"prepare", "prepare",
"Prepare data", "Prepare data",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "prepared"}, nil return agent.ToolResult{Content: "prepared"}, nil
}, },
) )
tool2 := agent.FunctionTool[Params]( require.NoError(t, err)
"finalize", tool2, err := agent.FunctionTool[Params](
"Finalize data", "finalize",
func(_ context.Context, _ Params) (agent.ToolResult, error) { "Finalize data",
t.Fatal("tool after handoff should not be executed") func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil t.Fatal("tool after handoff should not be executed")
}, return agent.ToolResult{}, nil
) },
)
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -2835,14 +2869,15 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
type Params struct{} type Params struct{}
failingTool := agent.FunctionTool[Params]( failingTool, err := agent.FunctionTool[Params](
"prepare", "prepare",
"Prepare data", "Prepare data",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, errors.New("preparation failed") return agent.ToolResult{}, errors.New("preparation failed")
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -2875,7 +2910,7 @@ func TestRun_HandoffWithPreHandoffTools(t *testing.T) {
agent.WithHandoffs(specialist), agent.WithHandoffs(specialist),
) )
_, err := router.Run( _, err = router.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("prepare and transfer")}, []llm.Message{userMessage("prepare and transfer")},
) )

2
pkg/agent/agent_tool.go
View File

@@ -40,7 +40,7 @@ type (
) )
var ( var (
agentToolParamsSchema = jsonSchemaFor[agentToolParams]() agentToolParamsSchema = mustJSONSchemaFor[agentToolParams]()
) )
func agentToolDepth(ctx context.Context) int { func agentToolDepth(ctx context.Context) int {

34
pkg/agent/agent_tool_test.go
View File

@@ -252,7 +252,7 @@ func TestAgentTool_Execute(t *testing.T) {
var captured string var captured string
type Params struct{} type Params struct{}
tenantTool := agent.FunctionTool[Params]( tenantTool, err := agent.FunctionTool[Params](
"get_tenant", "get_tenant",
"Get tenant", "Get tenant",
func(ctx context.Context, _ Params) (agent.ToolResult, error) { 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 return agent.ToolResult{Content: rc.TenantID}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -303,13 +304,14 @@ func TestAgentTool_Execute(t *testing.T) {
Expr string `json:"expr"` Expr string `json:"expr"`
} }
calcTool := agent.FunctionTool[Params]( calcTool, err := agent.FunctionTool[Params](
"calc", "calc",
"Calculate expression", "Calculate expression",
func(_ context.Context, p Params) (agent.ToolResult, error) { func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "42"}, nil return agent.ToolResult{Content: "42"}, nil
}, },
) )
require.NoError(t, err)
provider := &mockProvider{ provider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -379,13 +381,14 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
deleteTool := agent.FunctionTool[struct{}]( deleteTool, err := agent.FunctionTool[struct{}](
"delete_file", "delete_file",
"Delete a file", "Delete a file",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "file deleted"}, nil return agent.ToolResult{Content: "file deleted"}, nil
}, },
) )
require.NoError(t, err)
innerProvider := &mockProvider{ innerProvider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -422,7 +425,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")), agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")),
) )
_, err := outerAgent.Run( _, err = outerAgent.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("Delete the file")}, []llm.Message{userMessage("Delete the file")},
) )
@@ -443,7 +446,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
var toolExecuted bool var toolExecuted bool
deleteTool := agent.FunctionTool[struct{}]( deleteTool, err := agent.FunctionTool[struct{}](
"delete_file", "delete_file",
"Delete a file", "Delete a file",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { 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 return agent.ToolResult{Content: "file deleted"}, nil
}, },
) )
require.NoError(t, err)
innerProvider := &mockProvider{ innerProvider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -489,7 +493,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")), agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")),
) )
_, err := outerAgent.Run( _, err = outerAgent.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("Delete the file")}, []llm.Message{userMessage("Delete the file")},
) )
@@ -520,7 +524,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
func(t *testing.T) { func(t *testing.T) {
t.Parallel() t.Parallel()
deleteTool := agent.FunctionTool[struct{}]( deleteTool, err := agent.FunctionTool[struct{}](
"delete_file", "delete_file",
"Delete a file", "Delete a file",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
@@ -528,6 +532,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
innerProvider := &mockProvider{ innerProvider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -566,7 +571,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")), agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")),
) )
_, err := outerAgent.Run( _, err = outerAgent.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("Delete the file")}, []llm.Message{userMessage("Delete the file")},
) )
@@ -598,7 +603,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
var siblingCalled bool var siblingCalled bool
type Params struct{} type Params struct{}
siblingTool := agent.FunctionTool[Params]( siblingTool, err := agent.FunctionTool[Params](
"list_files", "list_files",
"List files", "List files",
func(_ context.Context, _ Params) (agent.ToolResult, error) { 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 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_file",
"Delete a file", "Delete a file",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
return agent.ToolResult{Content: "file deleted"}, nil return agent.ToolResult{Content: "file deleted"}, nil
}, },
) )
require.NoError(t, err)
innerProvider := &mockProvider{ innerProvider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -661,7 +668,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
), ),
) )
_, err := outerAgent.Run( _, err = outerAgent.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("List and delete files")}, []llm.Message{userMessage("List and delete files")},
) )
@@ -692,7 +699,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
var toolExecuted bool var toolExecuted bool
dangerTool := agent.FunctionTool[struct{}]( dangerTool, err := agent.FunctionTool[struct{}](
"danger", "danger",
"Dangerous operation", "Dangerous operation",
func(_ context.Context, _ struct{}) (agent.ToolResult, error) { 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 return agent.ToolResult{Content: "danger executed"}, nil
}, },
) )
require.NoError(t, err)
cProvider := &mockProvider{ cProvider := &mockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{
@@ -755,7 +763,7 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
agent.WithTools(agentB.AsTool("call_b", "Call agent B")), agent.WithTools(agentB.AsTool("call_b", "Call agent B")),
) )
_, err := agentA.Run( _, err = agentA.Run(
context.Background(), context.Background(),
[]llm.Message{userMessage("start")}, []llm.Message{userMessage("start")},
) )

2
pkg/agent/handoff.go
View File

@@ -45,7 +45,7 @@ type (
) )
var ( var (
handoffParamsSchema = jsonSchemaFor[handoffParams]() handoffParamsSchema = mustJSONSchemaFor[handoffParams]()
) )
func HandoffTo(agent *Agent, opts ...HandoffOption) *Handoff { func HandoffTo(agent *Agent, opts ...HandoffOption) *Handoff {

12
pkg/agent/output_type.go
View File

@@ -16,6 +16,7 @@ package agent
import ( import (
"encoding/json" "encoding/json"
"fmt"
"go.probo.inc/probo/pkg/llm" "go.probo.inc/probo/pkg/llm"
) )
@@ -27,11 +28,16 @@ type OutputType struct {
Schema json.RawMessage 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{ return &OutputType{
Name: name, Name: name,
Schema: jsonSchemaFor[T](), Schema: schema,
} }, nil
} }
func (o *OutputType) responseFormat() *llm.ResponseFormat { func (o *OutputType) responseFormat() *llm.ResponseFormat {

12
pkg/agent/output_type_test.go
View File

@@ -31,7 +31,8 @@ func TestNewOutputType_SetsNameAndSchema(t *testing.T) {
Score int `json:"score"` 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) assert.Equal(t, "test_result", ot.Name)
require.NotNil(t, ot.Schema) require.NotNil(t, ot.Schema)
@@ -50,7 +51,8 @@ func TestNewOutputType_EmptyStruct(t *testing.T) {
type Empty struct{} type Empty struct{}
ot := NewOutputType[Empty]("empty") ot, err := NewOutputType[Empty]("empty")
require.NoError(t, err)
assert.Equal(t, "empty", ot.Name) assert.Equal(t, "empty", ot.Name)
@@ -71,7 +73,8 @@ func TestOutputType_responseFormat(t *testing.T) {
Reason string `json:"reason"` Reason string `json:"reason"`
} }
ot := NewOutputType[Verdict]("verdict") ot, err := NewOutputType[Verdict]("verdict")
require.NoError(t, err)
rf := ot.responseFormat() rf := ot.responseFormat()
require.NotNil(t, rf) require.NotNil(t, rf)
@@ -91,7 +94,8 @@ func TestOutputType_responseFormat_SchemaMatchesOutputType(t *testing.T) {
Priority *int `json:"priority,omitempty"` Priority *int `json:"priority,omitempty"`
} }
ot := NewOutputType[Analysis]("analysis") ot, err := NewOutputType[Analysis]("analysis")
require.NoError(t, err)
rf := ot.responseFormat() rf := ot.responseFormat()
var schema map[string]any var schema map[string]any

16
pkg/agent/schema.go
View File

@@ -22,22 +22,30 @@ import (
"github.com/google/jsonschema-go/jsonschema" "github.com/google/jsonschema-go/jsonschema"
) )
func jsonSchemaFor[T any]() json.RawMessage { func jsonSchemaFor[T any]() (json.RawMessage, error) {
t := reflect.TypeFor[T]() t := reflect.TypeFor[T]()
schema, err := jsonschema.ForType(t, nil) schema, err := jsonschema.ForType(t, nil)
if err != 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) stripNullTypes(schema)
data, err := json.Marshal(schema) data, err := json.Marshal(schema)
if err != nil { 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 // stripNullTypes removes "null" from union types produced by pointer fields

33
pkg/agent/schema_test.go
View File

@@ -32,7 +32,8 @@ func TestGenerateSchema_PointerFieldsStripNull(t *testing.T) {
Done *bool `json:"done"` Done *bool `json:"done"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -65,7 +66,8 @@ func TestGenerateSchema_IntegerBoundsStripped(t *testing.T) {
Uint16 uint16 `json:"uint16"` Uint16 uint16 `json:"uint16"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -85,7 +87,8 @@ func TestGenerateSchema_EmptyStruct(t *testing.T) {
type Params struct{} type Params struct{}
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -104,7 +107,8 @@ func TestGenerateSchema_MapField(t *testing.T) {
Metadata map[string]string `json:"metadata"` Metadata map[string]string `json:"metadata"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -129,7 +133,8 @@ func TestGenerateSchema_NestedPointerStruct(t *testing.T) {
Inner *Inner `json:"inner"` Inner *Inner `json:"inner"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -153,7 +158,8 @@ func TestGenerateSchema_SliceOfPointers(t *testing.T) {
Names []*string `json:"names"` Names []*string `json:"names"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -175,7 +181,8 @@ func TestGenerateSchema_MapWithPointerValues(t *testing.T) {
Scores map[string]*int `json:"scores"` Scores map[string]*int `json:"scores"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) 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"` 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 var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -224,7 +232,8 @@ func TestGenerateSchema_RequiredVsOptional(t *testing.T) {
AlsoNeeded int `json:"also_needed"` AlsoNeeded int `json:"also_needed"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -252,7 +261,8 @@ func TestGenerateSchema_DeeplyNestedStructure(t *testing.T) {
Root Level1 `json:"root"` Root Level1 `json:"root"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))
@@ -293,7 +303,8 @@ func TestGenerateSchema_SliceOfStructs(t *testing.T) {
Items []Item `json:"items"` Items []Item `json:"items"`
} }
raw := jsonSchemaFor[Params]() raw, err := jsonSchemaFor[Params]()
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(raw, &schema)) require.NoError(t, json.Unmarshal(raw, &schema))

9
pkg/agent/tool.go
View File

@@ -52,8 +52,11 @@ func FunctionTool[P any](
name string, name string,
description string, description string,
fn func(ctx context.Context, params P) (ToolResult, error), fn func(ctx context.Context, params P) (ToolResult, error),
) Tool { ) (Tool, error) {
schema := jsonSchemaFor[P]() schema, err := jsonSchemaFor[P]()
if err != nil {
return nil, fmt.Errorf("cannot create tool %q: %w", name, err)
}
var parsed struct { var parsed struct {
Required []string `json:"required"` Required []string `json:"required"`
@@ -66,7 +69,7 @@ func FunctionTool[P any](
fn: fn, fn: fn,
schema: schema, schema: schema,
requiredFields: parsed.Required, requiredFields: parsed.Required,
} }, nil
} }
func (t *functionTool[P]) Name() string { return t.name } func (t *functionTool[P]) Name() string { return t.name }

58
pkg/agent/tool_test.go
View File

@@ -30,13 +30,14 @@ func TestFunctionTool_Name(t *testing.T) {
type Params struct{} type Params struct{}
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"my_tool", "my_tool",
"does things", "does things",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
assert.Equal(t, "my_tool", tool.Name()) assert.Equal(t, "my_tool", tool.Name())
} }
@@ -53,13 +54,14 @@ func TestFunctionTool_Definition(t *testing.T) {
Query string `json:"query"` Query string `json:"query"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"search", "search",
"Search for items", "Search for items",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
def := tool.Definition() def := tool.Definition()
assert.Equal(t, "search", def.Name) assert.Equal(t, "search", def.Name)
@@ -77,13 +79,14 @@ func TestFunctionTool_Definition(t *testing.T) {
Count int `json:"count"` Count int `json:"count"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"create", "create",
"Create items", "Create items",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
def := tool.Definition() def := tool.Definition()
require.NotNil(t, def.Parameters) require.NotNil(t, def.Parameters)
@@ -113,13 +116,14 @@ func TestFunctionTool_Definition(t *testing.T) {
type Params struct{} type Params struct{}
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"noop", "noop",
"No-op", "No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema)) require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema))
@@ -136,13 +140,14 @@ func TestFunctionTool_Definition(t *testing.T) {
Title *string `json:"title,omitempty"` Title *string `json:"title,omitempty"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"update", "update",
"Update", "Update",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
var schema map[string]any var schema map[string]any
require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema)) require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema))
@@ -168,13 +173,14 @@ func TestFunctionTool_Execute(t *testing.T) {
Y int `json:"y"` Y int `json:"y"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"add", "add",
"Add two numbers", "Add two numbers",
func(_ context.Context, p Params) (agent.ToolResult, error) { func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "42"}, nil return agent.ToolResult{Content: "42"}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{"x": 1, "y": 2}`) result, err := tool.Execute(context.Background(), `{"x": 1, "y": 2}`)
require.NoError(t, err) require.NoError(t, err)
@@ -193,7 +199,7 @@ func TestFunctionTool_Execute(t *testing.T) {
} }
var received string var received string
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"weather", "weather",
"Get weather", "Get weather",
func(_ context.Context, p Params) (agent.ToolResult, error) { func(_ context.Context, p Params) (agent.ToolResult, error) {
@@ -201,8 +207,9 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{Content: "sunny"}, nil 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) require.NoError(t, err)
assert.Equal(t, "Paris", received) assert.Equal(t, "Paris", received)
}, },
@@ -215,13 +222,14 @@ func TestFunctionTool_Execute(t *testing.T) {
type Params struct{} type Params struct{}
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"noop", "noop",
"No-op", "No-op",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "ok"}, nil return agent.ToolResult{Content: "ok"}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{invalid`) result, err := tool.Execute(context.Background(), `{invalid`)
require.NoError(t, err) require.NoError(t, err)
@@ -237,15 +245,16 @@ func TestFunctionTool_Execute(t *testing.T) {
type Params struct{} type Params struct{}
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"fail", "fail",
"Always fails", "Always fails",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, errors.New("db down") 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) require.Error(t, err)
assert.Contains(t, err.Error(), "db down") assert.Contains(t, err.Error(), "db down")
}, },
@@ -259,7 +268,7 @@ func TestFunctionTool_Execute(t *testing.T) {
type ctxKey struct{} type ctxKey struct{}
type Params struct{} type Params struct{}
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"ctx_check", "ctx_check",
"Check context", "Check context",
func(ctx context.Context, _ Params) (agent.ToolResult, error) { func(ctx context.Context, _ Params) (agent.ToolResult, error) {
@@ -267,6 +276,7 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{Content: val}, nil return agent.ToolResult{Content: val}, nil
}, },
) )
require.NoError(t, err)
ctx := context.WithValue(context.Background(), ctxKey{}, "hello") ctx := context.WithValue(context.Background(), ctxKey{}, "hello")
result, err := tool.Execute(ctx, `{}`) result, err := tool.Execute(ctx, `{}`)
@@ -284,7 +294,7 @@ func TestFunctionTool_Execute(t *testing.T) {
City string `json:"city"` City string `json:"city"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"weather", "weather",
"Get weather", "Get weather",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -292,6 +302,7 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{}`) result, err := tool.Execute(context.Background(), `{}`)
require.NoError(t, err) require.NoError(t, err)
@@ -310,7 +321,7 @@ func TestFunctionTool_Execute(t *testing.T) {
Country string `json:"country"` Country string `json:"country"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"weather", "weather",
"Get weather", "Get weather",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -318,6 +329,7 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{}`) result, err := tool.Execute(context.Background(), `{}`)
require.NoError(t, err) require.NoError(t, err)
@@ -337,7 +349,7 @@ func TestFunctionTool_Execute(t *testing.T) {
Country string `json:"country"` Country string `json:"country"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"weather", "weather",
"Get weather", "Get weather",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
@@ -345,6 +357,7 @@ func TestFunctionTool_Execute(t *testing.T) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{"city":"Paris"}`) result, err := tool.Execute(context.Background(), `{"city":"Paris"}`)
require.NoError(t, err) require.NoError(t, err)
@@ -363,13 +376,14 @@ func TestFunctionTool_Execute(t *testing.T) {
Units *string `json:"units,omitempty"` Units *string `json:"units,omitempty"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"weather", "weather",
"Get weather", "Get weather",
func(_ context.Context, p Params) (agent.ToolResult, error) { func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "sunny in " + p.City}, nil return agent.ToolResult{Content: "sunny in " + p.City}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{"city":"Paris"}`) result, err := tool.Execute(context.Background(), `{"city":"Paris"}`)
require.NoError(t, err) require.NoError(t, err)
@@ -387,13 +401,14 @@ func TestFunctionTool_Execute(t *testing.T) {
Name string `json:"name"` Name string `json:"name"`
} }
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"greet", "greet",
"Greet", "Greet",
func(_ context.Context, p Params) (agent.ToolResult, error) { func(_ context.Context, p Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "hi " + p.Name}, nil return agent.ToolResult{Content: "hi " + p.Name}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{"name":"Alice","extra":"ignored"}`) result, err := tool.Execute(context.Background(), `{"name":"Alice","extra":"ignored"}`)
require.NoError(t, err) require.NoError(t, err)
@@ -409,13 +424,14 @@ func TestFunctionTool_Execute(t *testing.T) {
type Params struct{} type Params struct{}
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"ping", "ping",
"Ping", "Ping",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "pong"}, nil return agent.ToolResult{Content: "pong"}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{}`) result, err := tool.Execute(context.Background(), `{}`)
require.NoError(t, err) require.NoError(t, err)
@@ -430,13 +446,14 @@ func TestFunctionTool_Execute(t *testing.T) {
type Params struct{} type Params struct{}
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"validate", "validate",
"Validate input", "Validate input",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{Content: "validation failed", IsError: true}, nil return agent.ToolResult{Content: "validation failed", IsError: true}, nil
}, },
) )
require.NoError(t, err)
result, err := tool.Execute(context.Background(), `{}`) result, err := tool.Execute(context.Background(), `{}`)
require.NoError(t, err) require.NoError(t, err)
@@ -451,13 +468,14 @@ func TestFunctionTool_InterfaceSatisfaction(t *testing.T) {
type Params struct{} type Params struct{}
tool := agent.FunctionTool( tool, err := agent.FunctionTool(
"test", "test",
"test tool", "test tool",
func(_ context.Context, _ Params) (agent.ToolResult, error) { func(_ context.Context, _ Params) (agent.ToolResult, error) {
return agent.ToolResult{}, nil return agent.ToolResult{}, nil
}, },
) )
require.NoError(t, err)
assert.Implements(t, (*agent.Tool)(nil), tool) assert.Implements(t, (*agent.Tool)(nil), tool)
assert.Implements(t, (*agent.ToolDescriptor)(nil), tool) assert.Implements(t, (*agent.ToolDescriptor)(nil), tool)

7
pkg/agent/typed.go
View File

@@ -36,11 +36,16 @@ func RunTyped[T any](
) (*TypedResult[T], error) { ) (*TypedResult[T], error) {
typed := a.clone() typed := a.clone()
schema, err := jsonSchemaFor[T]()
if err != nil {
return nil, fmt.Errorf("cannot create typed runner: %w", err)
}
typed.responseFormat = &llm.ResponseFormat{ typed.responseFormat = &llm.ResponseFormat{
Type: llm.ResponseFormatJSONSchema, Type: llm.ResponseFormatJSONSchema,
JSONSchema: &llm.JSONSchema{ JSONSchema: &llm.JSONSchema{
Name: typeName[T](), Name: typeName[T](),
Schema: jsonSchemaFor[T](), Schema: schema,
Strict: true, Strict: true,
}, },
} }

3
pkg/agent/typed_test.go
View File

@@ -404,13 +404,14 @@ func TestRunTyped(t *testing.T) {
City string `json:"city"` City string `json:"city"`
} }
weatherTool := FunctionTool[Params]( weatherTool, err := FunctionTool[Params](
"get_weather", "get_weather",
"Get weather for a city", "Get weather for a city",
func(_ context.Context, p Params) (ToolResult, error) { func(_ context.Context, p Params) (ToolResult, error) {
return ToolResult{Content: "Sunny, 22°C in " + p.City}, nil return ToolResult{Content: "Sunny, 22°C in " + p.City}, nil
}, },
) )
require.NoError(t, err)
provider := &typedMockProvider{ provider := &typedMockProvider{
responses: []*llm.ChatCompletionResponse{ responses: []*llm.ChatCompletionResponse{