Add agent library

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
2026-03-13 14:30:47 +01:00
parent 158c36d9ab
commit 9fd251ee5d
40 changed files with 9231 additions and 38 deletions
--- a/pkg/agent/tool_test.go
+++ b/pkg/agent/tool_test.go
@@ -0,0 +1,361 @@
+// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package agent_test
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.probo.inc/probo/pkg/agent"
+)
+
+func TestFunctionTool_Name(t *testing.T) {
+	t.Parallel()
+
+	type Params struct{}
+
+	tool := agent.FunctionTool(
+		"my_tool",
+		"does things",
+		func(_ context.Context, _ Params) (agent.ToolResult, error) {
+			return agent.ToolResult{}, nil
+		},
+	)
+
+	assert.Equal(t, "my_tool", tool.Name())
+}
+
+func TestFunctionTool_Definition(t *testing.T) {
+	t.Parallel()
+
+	t.Run(
+		"returns name and description",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct {
+				Query string `json:"query"`
+			}
+
+			tool := agent.FunctionTool(
+				"search",
+				"Search for items",
+				func(_ context.Context, _ Params) (agent.ToolResult, error) {
+					return agent.ToolResult{}, nil
+				},
+			)
+
+			def := tool.Definition()
+			assert.Equal(t, "search", def.Name)
+			assert.Equal(t, "Search for items", def.Description)
+		},
+	)
+
+	t.Run(
+		"generates valid JSON schema from params type",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct {
+				Name  string `json:"name" jsonschema:"The item name"`
+				Count int    `json:"count"`
+			}
+
+			tool := agent.FunctionTool(
+				"create",
+				"Create items",
+				func(_ context.Context, _ Params) (agent.ToolResult, error) {
+					return agent.ToolResult{}, nil
+				},
+			)
+
+			def := tool.Definition()
+			require.NotNil(t, def.Parameters)
+
+			var schema map[string]any
+			require.NoError(t, json.Unmarshal(def.Parameters, &schema))
+
+			assert.Equal(t, "object", schema["type"])
+
+			props := schema["properties"].(map[string]any)
+			assert.Contains(t, props, "name")
+			assert.Contains(t, props, "count")
+
+			nameProp := props["name"].(map[string]any)
+			assert.Equal(t, "string", nameProp["type"])
+			assert.Equal(t, "The item name", nameProp["description"])
+
+			countProp := props["count"].(map[string]any)
+			assert.Equal(t, "integer", countProp["type"])
+		},
+	)
+
+	t.Run(
+		"empty struct produces object schema with no properties",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct{}
+
+			tool := agent.FunctionTool(
+				"noop",
+				"No-op",
+				func(_ context.Context, _ Params) (agent.ToolResult, error) {
+					return agent.ToolResult{}, nil
+				},
+			)
+
+			var schema map[string]any
+			require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema))
+			assert.Equal(t, "object", schema["type"])
+		},
+	)
+
+	t.Run(
+		"pointer fields are not nullable in schema",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct {
+				Title *string `json:"title,omitempty"`
+			}
+
+			tool := agent.FunctionTool(
+				"update",
+				"Update",
+				func(_ context.Context, _ Params) (agent.ToolResult, error) {
+					return agent.ToolResult{}, nil
+				},
+			)
+
+			var schema map[string]any
+			require.NoError(t, json.Unmarshal(tool.Definition().Parameters, &schema))
+
+			props := schema["properties"].(map[string]any)
+			titleProp := props["title"].(map[string]any)
+			assert.Equal(t, "string", titleProp["type"])
+			assert.Nil(t, titleProp["types"])
+		},
+	)
+}
+
+func TestFunctionTool_Execute(t *testing.T) {
+	t.Parallel()
+
+	t.Run(
+		"unmarshals params and calls function",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct {
+				X int `json:"x"`
+				Y int `json:"y"`
+			}
+
+			tool := agent.FunctionTool(
+				"add",
+				"Add two numbers",
+				func(_ context.Context, p Params) (agent.ToolResult, error) {
+					return agent.ToolResult{Content: "42"}, nil
+				},
+			)
+
+			result, err := tool.Execute(context.Background(), `{"x": 1, "y": 2}`)
+			require.NoError(t, err)
+			assert.Equal(t, "42", result.Content)
+			assert.False(t, result.IsError)
+		},
+	)
+
+	t.Run(
+		"passes received params to function",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct {
+				City string `json:"city"`
+			}
+
+			var received string
+			tool := agent.FunctionTool(
+				"weather",
+				"Get weather",
+				func(_ context.Context, p Params) (agent.ToolResult, error) {
+					received = p.City
+					return agent.ToolResult{Content: "sunny"}, nil
+				},
+			)
+
+			_, err := tool.Execute(context.Background(), `{"city":"Paris"}`)
+			require.NoError(t, err)
+			assert.Equal(t, "Paris", received)
+		},
+	)
+
+	t.Run(
+		"invalid JSON returns tool error not Go error",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct{}
+
+			tool := agent.FunctionTool(
+				"noop",
+				"No-op",
+				func(_ context.Context, _ Params) (agent.ToolResult, error) {
+					return agent.ToolResult{Content: "ok"}, nil
+				},
+			)
+
+			result, err := tool.Execute(context.Background(), `{invalid`)
+			require.NoError(t, err)
+			assert.True(t, result.IsError)
+			assert.Contains(t, result.Content, "Invalid parameters")
+		},
+	)
+
+	t.Run(
+		"infrastructure error propagated as Go error",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct{}
+
+			tool := agent.FunctionTool(
+				"fail",
+				"Always fails",
+				func(_ context.Context, _ Params) (agent.ToolResult, error) {
+					return agent.ToolResult{}, errors.New("db down")
+				},
+			)
+
+			_, err := tool.Execute(context.Background(), `{}`)
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "db down")
+		},
+	)
+
+	t.Run(
+		"context is forwarded to function",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type ctxKey struct{}
+			type Params struct{}
+
+			tool := agent.FunctionTool(
+				"ctx_check",
+				"Check context",
+				func(ctx context.Context, _ Params) (agent.ToolResult, error) {
+					val := ctx.Value(ctxKey{}).(string)
+					return agent.ToolResult{Content: val}, nil
+				},
+			)
+
+			ctx := context.WithValue(context.Background(), ctxKey{}, "hello")
+			result, err := tool.Execute(ctx, `{}`)
+			require.NoError(t, err)
+			assert.Equal(t, "hello", result.Content)
+		},
+	)
+
+	t.Run(
+		"extra JSON fields are ignored",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct {
+				Name string `json:"name"`
+			}
+
+			tool := agent.FunctionTool(
+				"greet",
+				"Greet",
+				func(_ context.Context, p Params) (agent.ToolResult, error) {
+					return agent.ToolResult{Content: "hi " + p.Name}, nil
+				},
+			)
+
+			result, err := tool.Execute(context.Background(), `{"name":"Alice","extra":"ignored"}`)
+			require.NoError(t, err)
+			assert.Equal(t, "hi Alice", result.Content)
+			assert.False(t, result.IsError)
+		},
+	)
+
+	t.Run(
+		"empty JSON object works for empty params",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct{}
+
+			tool := agent.FunctionTool(
+				"ping",
+				"Ping",
+				func(_ context.Context, _ Params) (agent.ToolResult, error) {
+					return agent.ToolResult{Content: "pong"}, nil
+				},
+			)
+
+			result, err := tool.Execute(context.Background(), `{}`)
+			require.NoError(t, err)
+			assert.Equal(t, "pong", result.Content)
+		},
+	)
+
+	t.Run(
+		"function can return IsError true",
+		func(t *testing.T) {
+			t.Parallel()
+
+			type Params struct{}
+
+			tool := agent.FunctionTool(
+				"validate",
+				"Validate input",
+				func(_ context.Context, _ Params) (agent.ToolResult, error) {
+					return agent.ToolResult{Content: "validation failed", IsError: true}, nil
+				},
+			)
+
+			result, err := tool.Execute(context.Background(), `{}`)
+			require.NoError(t, err)
+			assert.True(t, result.IsError)
+			assert.Equal(t, "validation failed", result.Content)
+		},
+	)
+}
+
+func TestFunctionTool_InterfaceSatisfaction(t *testing.T) {
+	t.Parallel()
+
+	type Params struct{}
+
+	tool := agent.FunctionTool(
+		"test",
+		"test tool",
+		func(_ context.Context, _ Params) (agent.ToolResult, error) {
+			return agent.ToolResult{}, nil
+		},
+	)
+
+	assert.Implements(t, (*agent.Tool)(nil), tool)
+	assert.Implements(t, (*agent.ToolDescriptor)(nil), tool)
+}