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