summaryrefslogtreecommitdiff
path: root/examples/simple-agent-go/main.go
blob: 8a02142d2926a153840ba527291d610e9ccae09f (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
// A tiny standalone libpanto-go app: one prompt in, streamed text out.
//
// Run with:
//
//	PANTO_API_STYLE=openai_chat|anthropic_messages \
//	PANTO_API_KEY=... \
//	PANTO_BASE_URL=... \
//	PANTO_MODEL=... \
//	go run examples/simple-agent-go/main.go "list the files"

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"strings"

	panto "code.tjp.lol/libpantograph.git/libpanto-go"
)

func main() {
	panto.Init()
	defer panto.Deinit()

	apiStyle := env("PANTO_API_STYLE", "openai_chat")
	apiKey := requireEnv("PANTO_API_KEY")
	baseURL := requireEnv("PANTO_BASE_URL")
	model := requireEnv("PANTO_MODEL")
	maxTokens := uint32(4096)
	if v := os.Getenv("PANTO_MAX_TOKENS"); v != "" {
		fmt.Sscanf(v, "%d", &maxTokens)
	}

	prompt := strings.Join(os.Args[1:], " ")
	if prompt == "" {
		fmt.Fprintln(os.Stderr, "usage: simple-agent-go <prompt>")
		os.Exit(2)
	}

	var cfg panto.Config
	if strings.HasPrefix(apiStyle, "anthropic") {
		cfg.Provider = panto.ProviderConfig{
			AnthropicMessages: &panto.AnthropicMessagesConfig{
				APIKey:    apiKey,
				BaseURL:   baseURL,
				Model:     model,
				MaxTokens: maxTokens,
			},
		}
	} else {
		cfg.Provider = panto.ProviderConfig{
			OpenAIChat: &panto.OpenAIChatConfig{
				APIKey:    apiKey,
				BaseURL:   baseURL,
				Model:     model,
				MaxTokens: maxTokens,
			},
		}
	}

	agent, err := panto.NewAgent(cfg)
	if err != nil {
		panic(err)
	}
	defer agent.Close()

	agent.SetSystemPrompt("You are a very small coding agent.\nUse the bash tool when you need to inspect or change the local workspace.\nKeep final answers concise.")

	agent.RegisterTool(panto.Tool{
		Decl: panto.ToolDecl{
			Name:        "bash",
			Description: "Run a shell command with bash -lc and return combined stdout/stderr plus the exit code.",
			SchemaJSON: `{
				"type": "object",
				"properties": {
					"command": {
						"type": "string",
						"description": "The command to run."
					}
				},
				"required": ["command"]
			}`,
		},
		Invoke: func(input string) (panto.ResultParts, error) {
			var params map[string]any
			if err := json.Unmarshal([]byte(input), &params); err != nil {
				return nil, fmt.Errorf("invalid input: %w", err)
			}
			command, ok := params["command"].(string)
			if !ok || command == "" {
				return nil, fmt.Errorf("bash.command must be a non-empty string")
			}

			out, err := runBash(command)
			if err != nil {
				return nil, err
			}
			return panto.ResultParts{{Text: out}}, nil
		},
	})

	stream, err := agent.Run(panto.UserMessage(panto.TextBlock(prompt)))
	if err != nil {
		panic(err)
	}
	defer stream.Close()

	for ev := range stream.Iter() {
		switch ev := ev.(type) {
		case panto.ContentDeltaEvent:
			fmt.Print(ev.Delta)
			os.Stdout.Sync()
		case panto.ToolDispatchStartEvent:
			fmt.Fprintf(os.Stderr, "\n[bash: running %d tool call(s)]\n", ev.Count)
		case panto.ToolDispatchCompleteEvent:
			fmt.Fprint(os.Stderr, "[bash: done]\n")
		}
	}
	if err := stream.Err(); err != nil {
		panic(err)
	}
	fmt.Println()
}

func env(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

func requireEnv(key string) string {
	v := os.Getenv(key)
	if v == "" {
		panic("missing required env var: " + key)
	}
	return v
}

func runBash(command string) (string, error) {
	cmd := exec.Command("bash", "-lc", command)
	out, err := cmd.CombinedOutput()
	status := "unknown"
	if exitErr, ok := err.(*exec.ExitError); ok {
		status = fmt.Sprintf("%d", exitErr.ExitCode())
	} else if err != nil {
		return "", err
	}
	return fmt.Sprintf("exit_code=%s\n%s", status, string(out)), nil
}