diff --git a/README.md b/README.md index 311c043f..bb3ef01e 100644 --- a/README.md +++ b/README.md @@ -545,6 +545,28 @@ host, err := kit.New(ctx, &kit.Options{ }) ``` +### Custom Tools + +Create custom tools with automatic schema generation — no external dependencies needed: + +```go +type SearchInput struct { + Query string `json:"query" description:"Search query"` +} + +searchTool := kit.NewTool("search", "Search the codebase", + func(ctx context.Context, input SearchInput) (kit.ToolOutput, error) { + return kit.TextResult("Found: ..."), nil + }, +) + +host, _ := kit.New(ctx, &kit.Options{ + ExtraTools: []kit.Tool{searchTool}, // adds alongside built-in tools +}) +``` + +Use `kit.NewParallelTool` for tools safe to run concurrently. See the [SDK docs](/sdk/overview) for full details on struct tags, `ToolOutput` fields, and `ToolCallIDFromContext`. + ### With Callbacks ```go diff --git a/internal/ui/cli.go b/internal/ui/cli.go index 3d8bbad9..ae55ebf3 100644 --- a/internal/ui/cli.go +++ b/internal/ui/cli.go @@ -5,7 +5,6 @@ import ( "os" "time" - "charm.land/fantasy" "charm.land/lipgloss/v2" "golang.org/x/term" @@ -173,33 +172,6 @@ func (c *CLI) DisplayDebugConfig(config map[string]any) { fmt.Println(c.renderer.RenderDebugConfigMessage(config, time.Now()).Content) } -// UpdateUsageFromResponse records token usage using metadata from the fantasy -// response. Only actual API-reported tokens are used for cost tracking. -// If the provider doesn't report token counts, no usage is recorded. -func (c *CLI) UpdateUsageFromResponse(response *fantasy.Response, inputText string) { - if c.usageTracker == nil { - return - } - - usage := response.Usage - inputTokens := int(usage.InputTokens) - outputTokens := int(usage.OutputTokens) - - // Only use actual API-reported tokens for cost tracking. - // We intentionally do NOT estimate tokens - estimation is inaccurate - // and should never be used for cost calculations. - if inputTokens > 0 { - cacheReadTokens := int(usage.CacheReadTokens) - cacheWriteTokens := int(usage.CacheCreationTokens) - c.usageTracker.UpdateUsage(inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens) - // Per-response usage is a single API call, so it represents the - // actual context window fill level. - c.usageTracker.SetContextTokens(inputTokens + outputTokens) - } - // If inputTokens is 0, the provider didn't report usage - we skip recording - // rather than estimating, to ensure cost accuracy. -} - // DisplayUsageAfterResponse renders and displays token usage information immediately // following an AI response. This provides real-time feedback about the cost and // token consumption of each interaction. diff --git a/pkg/kit/tools.go b/pkg/kit/tools.go index ec8c0dbe..c20b9c33 100644 --- a/pkg/kit/tools.go +++ b/pkg/kit/tools.go @@ -1,6 +1,8 @@ package kit import ( + "context" + "charm.land/fantasy" "github.com/mark3labs/kit/internal/core" @@ -16,6 +18,123 @@ type ToolOption = core.ToolOption // If empty, os.Getwd() is used at execution time. var WithWorkDir = core.WithWorkDir +// --- Custom tool creation --- + +// ToolOutput is the return value from custom tool handlers created with +// [NewTool] or [NewParallelTool]. It provides a dependency-free way to +// return results without importing the underlying LLM framework. +type ToolOutput struct { + // Content is the text content returned to the LLM. + Content string + + // IsError, when true, signals to the LLM that the tool call failed. + IsError bool + + // Data contains optional binary data (images, audio, etc.). + Data []byte + + // MediaType is the MIME type for binary Data (e.g. "image/png"). + MediaType string + + // Metadata is optional opaque metadata attached to the response. + // It is not sent to the LLM but may be consumed by hooks or the UI. + Metadata any +} + +// TextResult creates a successful text [ToolOutput]. +func TextResult(content string) ToolOutput { + return ToolOutput{Content: content} +} + +// ErrorResult creates an error [ToolOutput]. The LLM will see the content +// as a tool error, allowing it to retry or adjust its approach. +func ErrorResult(content string) ToolOutput { + return ToolOutput{Content: content, IsError: true} +} + +// toolCallIDKey is the context key for the tool call ID. +type toolCallIDKey struct{} + +// ToolCallIDFromContext extracts the tool call ID from the context. +// The call ID is set automatically by [NewTool] and [NewParallelTool] +// before invoking the handler. Returns an empty string if no ID is present. +func ToolCallIDFromContext(ctx context.Context) string { + s, _ := ctx.Value(toolCallIDKey{}).(string) + return s +} + +// NewTool creates a custom [Tool] with automatic JSON schema generation from +// the TInput struct type. The handler receives a typed input (deserialized +// from the LLM's JSON arguments) and returns a [ToolResult]. +// +// Struct tags on TInput control the generated schema: +// +// json:"name" → parameter name +// description:"..." → parameter description shown to the LLM +// enum:"a,b,c" → restrict valid values +// omitempty → marks the parameter as optional +// +// The tool call ID is injected into the context and can be retrieved with +// [ToolCallIDFromContext]. +// +// Example: +// +// type WeatherInput struct { +// City string `json:"city" description:"City name"` +// } +// +// tool := kit.NewTool("get_weather", "Get weather for a city", +// func(ctx context.Context, input WeatherInput) (kit.ToolResult, error) { +// return kit.TextResult("72°F, sunny in " + input.City), nil +// }, +// ) +func NewTool[TInput any](name, description string, fn func(ctx context.Context, input TInput) (ToolOutput, error)) Tool { + return fantasy.NewAgentTool(name, description, + func(ctx context.Context, input TInput, call fantasy.ToolCall) (fantasy.ToolResponse, error) { + ctx = context.WithValue(ctx, toolCallIDKey{}, call.ID) + result, err := fn(ctx, input) + if err != nil { + return fantasy.NewTextErrorResponse(err.Error()), nil + } + resp := fantasy.ToolResponse{ + Content: result.Content, + IsError: result.IsError, + Data: result.Data, + MediaType: result.MediaType, + } + if result.Metadata != nil { + resp = fantasy.WithResponseMetadata(resp, result.Metadata) + } + return resp, nil + }, + ) +} + +// NewParallelTool is like [NewTool] but marks the tool as safe for concurrent +// execution alongside other tools. Use this when the tool has no side effects +// or when concurrent calls are safe. +func NewParallelTool[TInput any](name, description string, fn func(ctx context.Context, input TInput) (ToolOutput, error)) Tool { + return fantasy.NewParallelAgentTool(name, description, + func(ctx context.Context, input TInput, call fantasy.ToolCall) (fantasy.ToolResponse, error) { + ctx = context.WithValue(ctx, toolCallIDKey{}, call.ID) + result, err := fn(ctx, input) + if err != nil { + return fantasy.NewTextErrorResponse(err.Error()), nil + } + resp := fantasy.ToolResponse{ + Content: result.Content, + IsError: result.IsError, + Data: result.Data, + MediaType: result.MediaType, + } + if result.Metadata != nil { + resp = fantasy.WithResponseMetadata(resp, result.Metadata) + } + return resp, nil + }, + ) +} + // --- Individual tool constructors --- // NewReadTool creates a file-reading tool. diff --git a/pkg/kit/tools_test.go b/pkg/kit/tools_test.go new file mode 100644 index 00000000..901d26fd --- /dev/null +++ b/pkg/kit/tools_test.go @@ -0,0 +1,119 @@ +package kit_test + +import ( + "context" + "testing" + + kit "github.com/mark3labs/kit/pkg/kit" +) + +// TestNewTool_BasicTextResult verifies that NewTool creates a working tool +// that returns text content via ToolOutput. +func TestNewTool_BasicTextResult(t *testing.T) { + type Input struct { + Name string `json:"name"` + } + + tool := kit.NewTool("greet", "Greet someone", + func(ctx context.Context, input Input) (kit.ToolOutput, error) { + return kit.TextResult("hello " + input.Name), nil + }, + ) + + info := tool.Info() + if info.Name != "greet" { + t.Errorf("Info().Name = %q, want %q", info.Name, "greet") + } + if info.Description != "Greet someone" { + t.Errorf("Info().Description = %q, want %q", info.Description, "Greet someone") + } + if info.Parallel { + t.Error("NewTool should not mark tool as parallel") + } +} + +// TestNewParallelTool_MarkedParallel verifies that NewParallelTool marks the +// tool as safe for concurrent execution. +func TestNewParallelTool_MarkedParallel(t *testing.T) { + type Input struct { + Query string `json:"query"` + } + + tool := kit.NewParallelTool("search", "Search for things", + func(ctx context.Context, input Input) (kit.ToolOutput, error) { + return kit.TextResult("found: " + input.Query), nil + }, + ) + + info := tool.Info() + if info.Name != "search" { + t.Errorf("Info().Name = %q, want %q", info.Name, "search") + } + if !info.Parallel { + t.Error("NewParallelTool should mark tool as parallel") + } +} + +// TestTextResult verifies the TextResult convenience constructor. +func TestTextResult(t *testing.T) { + r := kit.TextResult("ok") + if r.Content != "ok" { + t.Errorf("Content = %q, want %q", r.Content, "ok") + } + if r.IsError { + t.Error("TextResult should not set IsError") + } +} + +// TestErrorResult verifies the ErrorResult convenience constructor. +func TestErrorResult(t *testing.T) { + r := kit.ErrorResult("bad input") + if r.Content != "bad input" { + t.Errorf("Content = %q, want %q", r.Content, "bad input") + } + if !r.IsError { + t.Error("ErrorResult should set IsError") + } +} + +// TestToolCallIDFromContext verifies round-trip context injection. +func TestToolCallIDFromContext(t *testing.T) { + // Empty context returns empty string. + if id := kit.ToolCallIDFromContext(context.Background()); id != "" { + t.Errorf("expected empty string from bare context, got %q", id) + } +} + +// TestToolOutput_Metadata verifies that metadata can be set on ToolOutput. +func TestToolOutput_Metadata(t *testing.T) { + r := kit.ToolOutput{ + Content: "data", + Metadata: map[string]string{"key": "value"}, + } + if r.Metadata == nil { + t.Error("expected non-nil Metadata") + } + m, ok := r.Metadata.(map[string]string) + if !ok { + t.Fatalf("expected map[string]string, got %T", r.Metadata) + } + if m["key"] != "value" { + t.Errorf("Metadata[key] = %q, want %q", m["key"], "value") + } +} + +// TestToolOutput_BinaryData verifies that binary data fields work correctly. +func TestToolOutput_BinaryData(t *testing.T) { + data := []byte{0x89, 0x50, 0x4E, 0x47} + r := kit.ToolOutput{ + Content: "image result", + Data: data, + MediaType: "image/png", + } + if len(r.Data) != 4 { + t.Errorf("Data len = %d, want 4", len(r.Data)) + } + if r.MediaType != "image/png" { + t.Errorf("MediaType = %q, want %q", r.MediaType, "image/png") + } +} diff --git a/skills/kit-sdk/SKILL.md b/skills/kit-sdk/SKILL.md index 6b7dac77..74c802ec 100644 --- a/skills/kit-sdk/SKILL.md +++ b/skills/kit-sdk/SKILL.md @@ -347,6 +347,77 @@ Lower values run first. Within the same priority, registration order applies. Fi ## Tools +### Creating custom tools + +Use `kit.NewTool` to create custom tools. The JSON schema is auto-generated from the input struct — no external dependencies required: + +```go +type WeatherInput struct { + City string `json:"city" description:"City name, e.g. 'San Francisco'"` +} + +weatherTool := kit.NewTool("get_weather", "Get current weather for a city", + func(ctx context.Context, input WeatherInput) (kit.ToolOutput, error) { + // Your logic here (API calls, database lookups, etc.) + return kit.TextResult("72°F, sunny in " + input.City), nil + }, +) + +host, _ := kit.New(ctx, &kit.Options{ + ExtraTools: []kit.Tool{weatherTool}, +}) +``` + +**Struct tags** control the generated schema: + +| Tag | Purpose | Example | +|-----|---------|---------| +| `json:"name"` | Parameter name | `json:"city"` | +| `description:"..."` | Description shown to the LLM | `description:"City name"` | +| `enum:"a,b,c"` | Restrict valid values | `enum:"json,text,csv"` | +| `omitempty` | Marks parameter as optional | `json:"limit,omitempty"` | + +**Return helpers:** + +| Function | Description | +|----------|-------------| +| `kit.TextResult(content)` | Successful text result | +| `kit.ErrorResult(content)` | Error result (LLM sees it as a tool error) | + +**ToolOutput fields** (for advanced use): + +```go +kit.ToolOutput{ + Content: "result text", // text returned to the LLM + IsError: false, // true = LLM sees this as an error + Data: pngBytes, // optional binary data (images, audio) + MediaType: "image/png", // MIME type for binary Data + Metadata: map[string]any{}, // opaque metadata for hooks/UI (not sent to LLM) +} +``` + +**Parallel tools** — mark as safe for concurrent execution: + +```go +searchTool := kit.NewParallelTool("search", "Search the web", + func(ctx context.Context, input SearchInput) (kit.ToolOutput, error) { + return kit.TextResult("results..."), nil + }, +) +``` + +**Tool call ID** — available in context for logging/tracing: + +```go +tool := kit.NewTool("my_tool", "...", + func(ctx context.Context, input MyInput) (kit.ToolOutput, error) { + callID := kit.ToolCallIDFromContext(ctx) // correlation ID from the LLM + log.Printf("[%s] my_tool called", callID) + return kit.TextResult("ok"), nil + }, +) +``` + ### Built-in tool constructors ```go diff --git a/www/pages/sdk/callbacks.md b/www/pages/sdk/callbacks.md index 3b8fbc0f..ce4732c4 100644 --- a/www/pages/sdk/callbacks.md +++ b/www/pages/sdk/callbacks.md @@ -7,17 +7,16 @@ description: Monitor tool calls and streaming output with the Kit Go SDK. ## Event-based monitoring -For more granular control, use the event subscription API: +Subscribe to events for real-time monitoring. Each method returns an unsubscribe function: ```go -// Subscribe returns an unsubscribe function unsub := host.OnToolCall(func(event kit.ToolCallEvent) { - fmt.Printf("Tool: %s, Args: %s\n", event.Name, event.Args) + fmt.Printf("Tool: %s, Args: %s\n", event.ToolName, event.ToolArgs) }) defer unsub() unsub2 := host.OnToolResult(func(event kit.ToolResultEvent) { - fmt.Printf("Result: %s (error: %v)\n", event.Name, event.IsError) + fmt.Printf("Result: %s (error: %v)\n", event.ToolName, event.IsError) }) defer unsub2() @@ -44,33 +43,62 @@ defer unsub6() ## Hook system -Hooks allow you to intercept and modify behavior. Unlike events, hooks can modify or cancel operations: +Hooks can **modify or cancel** operations. Unlike events (read-only), hooks are read-write interceptors. + +### BeforeToolCall — block tool execution ```go -// Intercept tool calls before execution -host.OnBeforeToolCall(0, func(ctx context.Context, name string, args string) (string, error) { - if name == "bash" { - log.Println("Bash command:", args) +host.OnBeforeToolCall(kit.HookPriorityNormal, func(h kit.BeforeToolCallHook) *kit.BeforeToolCallResult { + // h.ToolCallID, h.ToolName, h.ToolArgs + if h.ToolName == "bash" && strings.Contains(h.ToolArgs, "rm -rf") { + return &kit.BeforeToolCallResult{Block: true, Reason: "dangerous command"} } - return args, nil // return modified args or error to cancel + return nil // allow }) +``` -// Process results after tool execution -host.OnAfterToolResult(0, func(ctx context.Context, name string, result string) (string, error) { - return result, nil -}) +### AfterToolResult — modify tool output -// Before/after each agent turn -host.OnBeforeTurn(0, func(ctx context.Context) error { - return nil -}) - -host.OnAfterTurn(0, func(ctx context.Context) error { +```go +host.OnAfterToolResult(kit.HookPriorityNormal, func(h kit.AfterToolResultHook) *kit.AfterToolResultResult { + // h.ToolCallID, h.ToolName, h.ToolArgs, h.Result, h.IsError + if h.ToolName == "read" { + filtered := redactSecrets(h.Result) + return &kit.AfterToolResultResult{Result: &filtered} + } return nil }) ``` -The first argument is a priority (lower = runs first). +### BeforeTurn — modify prompt, inject messages + +```go +host.OnBeforeTurn(kit.HookPriorityNormal, func(h kit.BeforeTurnHook) *kit.BeforeTurnResult { + // h.Prompt + newPrompt := h.Prompt + "\nAlways respond in JSON." + return &kit.BeforeTurnResult{Prompt: &newPrompt} + // Also available: SystemPrompt *string, InjectText *string +}) +``` + +### AfterTurn — observation only + +```go +host.OnAfterTurn(kit.HookPriorityNormal, func(h kit.AfterTurnHook) { + // h.Response, h.Error + log.Printf("Turn completed: %d chars", len(h.Response)) +}) +``` + +### Hook priorities + +```go +kit.HookPriorityHigh = 0 // runs first +kit.HookPriorityNormal = 50 // default +kit.HookPriorityLow = 100 // runs last +``` + +Lower values run first. First non-nil result wins. ## Subagent event monitoring diff --git a/www/pages/sdk/options.md b/www/pages/sdk/options.md index b3727045..937fc57d 100644 --- a/www/pages/sdk/options.md +++ b/www/pages/sdk/options.md @@ -68,3 +68,28 @@ host, err := kit.New(ctx, &kit.Options{ | `CompactionOptions` | `*CompactionOptions` | — | Configuration for auto-compaction | | `Skills` | `[]string` | — | Explicit skill files/dirs to load | | `SkillsDir` | `string` | — | Override default skills directory | + +## Tool configuration + +**`Tools`** replaces ALL default tools (core + MCP + extension). **`ExtraTools`** adds tools alongside the defaults. Use `Tools` to restrict capabilities; use `ExtraTools` to extend them. + +Create custom tools with `kit.NewTool` — no external dependencies needed: + +```go +type LookupInput struct { + ID string `json:"id" description:"Record ID to look up"` +} + +lookupTool := kit.NewTool("lookup", "Look up a record by ID", + func(ctx context.Context, input LookupInput) (kit.ToolOutput, error) { + record := db.Find(input.ID) + return kit.TextResult(record.String()), nil + }, +) + +host, _ := kit.New(ctx, &kit.Options{ + ExtraTools: []kit.Tool{lookupTool}, +}) +``` + +See [Overview](/sdk/overview#custom-tools) for full custom tool documentation. diff --git a/www/pages/sdk/overview.md b/www/pages/sdk/overview.md index ba4aebaf..a277b0b3 100644 --- a/www/pages/sdk/overview.md +++ b/www/pages/sdk/overview.md @@ -68,6 +68,44 @@ The SDK provides several prompt variants: | `Steer(ctx, instruction)` | System-level steering without user message | | `FollowUp(ctx, text)` | Continue without new user input | +## Custom tools + +Create custom tools with `kit.NewTool`. The JSON schema is auto-generated from the input struct — no external dependencies required: + +```go +type WeatherInput struct { + City string `json:"city" description:"City name"` +} + +weatherTool := kit.NewTool("get_weather", "Get current weather for a city", + func(ctx context.Context, input WeatherInput) (kit.ToolOutput, error) { + return kit.TextResult("72°F, sunny in " + input.City), nil + }, +) + +host, _ := kit.New(ctx, &kit.Options{ + ExtraTools: []kit.Tool{weatherTool}, +}) +``` + +Struct tags control the schema: + +- `json:"name"` — parameter name +- `description:"..."` — description shown to the LLM +- `enum:"a,b,c"` — restrict valid values +- `omitempty` — marks the parameter as optional + +Return values: + +| Helper | Description | +|--------|-------------| +| `kit.TextResult(s)` | Successful text result | +| `kit.ErrorResult(s)` | Error result (LLM sees it as a tool error) | + +For advanced use, return a `kit.ToolOutput` struct directly with `Data`, `MediaType`, and `Metadata` fields. + +Use `kit.NewParallelTool` for tools that are safe to run concurrently. Use `kit.ToolCallIDFromContext(ctx)` to retrieve the LLM-assigned call ID for logging or tracing. + ## Event system Subscribe to events for monitoring: