main.go

 1package main
 2
 3import (
 4	"context"
 5	"fmt"
 6	"os"
 7
 8	"github.com/charmbracelet/crush/internal/ai"
 9	"github.com/charmbracelet/crush/internal/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) {
62			fmt.Print(text)
63		},
64
65		// Show when tools are called
66		OnToolCall: func(toolCall ai.ToolCallContent) {
67			fmt.Printf("\n[Tool: %s called]\n", toolCall.ToolName)
68		},
69
70		// Show tool results
71		OnToolResult: func(result ai.ToolResultContent) {
72			fmt.Printf("[Tool result received]\n")
73		},
74
75		// Show when each step completes
76		OnStepFinish: func(step ai.StepResult) {
77			fmt.Printf("\n[Step completed: %s]\n", step.FinishReason)
78		},
79	}
80
81	fmt.Println("Assistant response:")
82	result, err := agent.Stream(ctx, streamCall)
83	if err != nil {
84		fmt.Printf("Error: %v\n", err)
85		os.Exit(1)
86	}
87
88	fmt.Printf("\n\nFinal result: %s\n", result.Response.Content.Text())
89	fmt.Printf("Steps: %d, Total tokens: %d\n", len(result.Steps), result.TotalUsage.TotalTokens)
90}