1package main
 2
 3import (
 4	"context"
 5	"fmt"
 6	"os"
 7
 8	"github.com/charmbracelet/ai"
 9	"github.com/charmbracelet/ai/providers"
10)
11
12func main() {
13	// Check for API key
14	apiKey := os.Getenv("OPENAI_API_KEY")
15	if apiKey == "" {
16		fmt.Println("Please set OPENAI_API_KEY environment variable")
17		os.Exit(1)
18	}
19
20	// Create provider and model
21	provider := providers.NewOpenAiProvider(
22		providers.WithOpenAiAPIKey(apiKey),
23	)
24	model, err := provider.LanguageModel("gpt-4o-mini")
25	if err != nil {
26		fmt.Println(err)
27		return
28	}
29
30	// Create echo tool using the new type-safe API
31	type EchoInput struct {
32		Message string `json:"message" description:"The message to echo back"`
33	}
34
35	echoTool := ai.NewAgentTool(
36		"echo",
37		"Echo back the provided message",
38		func(ctx context.Context, input EchoInput, _ ai.ToolCall) (ai.ToolResponse, error) {
39			return ai.NewTextResponse("Echo: " + input.Message), nil
40		},
41	)
42
43	// Create streaming agent
44	agent := ai.NewAgent(
45		model,
46		ai.WithSystemPrompt("You are a helpful assistant."),
47		ai.WithTools(echoTool),
48	)
49
50	ctx := context.Background()
51
52	fmt.Println("Simple Streaming Agent Example")
53	fmt.Println("==============================")
54	fmt.Println()
55
56	// Basic streaming with key callbacks
57	streamCall := ai.AgentStreamCall{
58		Prompt: "Please echo back 'Hello, streaming world!'",
59
60		// Show real-time text as it streams
61		OnTextDelta: func(id, text string) error {
62			fmt.Print(text)
63			return nil
64		},
65
66		// Show when tools are called
67		OnToolCall: func(toolCall ai.ToolCallContent) error {
68			fmt.Printf("\n[Tool: %s called]\n", toolCall.ToolName)
69			return nil
70		},
71
72		// Show tool results
73		OnToolResult: func(result ai.ToolResultContent) error {
74			fmt.Printf("[Tool result received]\n")
75			return nil
76		},
77
78		// Show when each step completes
79		OnStepFinish: func(step ai.StepResult) error {
80			fmt.Printf("\n[Step completed: %s]\n", step.FinishReason)
81			return nil
82		},
83	}
84
85	fmt.Println("Assistant response:")
86	result, err := agent.Stream(ctx, streamCall)
87	if err != nil {
88		fmt.Printf("Error: %v\n", err)
89		os.Exit(1)
90	}
91
92	fmt.Printf("\n\nFinal result: %s\n", result.Response.Content.Text())
93	fmt.Printf("Steps: %d, Total tokens: %d\n", len(result.Steps), result.TotalUsage.TotalTokens)
94}