A Tool-Calling Loop

{
  "type": "function",
  "function": {
    "name": "get_gpu_status",
    "description": "Return memory use and temperature for one GPU on one node.",
    "parameters": {
      "type": "object",
      "properties": {
        "node": {"type": "string"},
        "gpu":  {"type": "integer", "minimum": 0}
      },
      "required": ["node", "gpu"]
    }
  }
}

// agent.go — the tool-calling loop: registry, validation, limits, and error feedback.
package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"
)

type Message struct {
	Role    string // system, user, assistant, tool
	Content string
	Call    *ToolCall // set when the assistant asks for a tool
}

type ToolCall struct {
	Name string
	Args json.RawMessage
}

type Tool struct {
	Name, Description string
	ReadOnly          bool
	Run               func(ctx context.Context, args json.RawMessage) (string, error)
}

// decode parses arguments strictly: unknown fields are an error.
func decode(args json.RawMessage, into any) error {
	dec := json.NewDecoder(strings.NewReader(string(args)))
	dec.DisallowUnknownFields()
	return dec.Decode(into)
}

var fleet = map[string][]int{"n7": {71, 68, 93, 70}} // node → GPU temperatures

func tools() map[string]Tool {
	return map[string]Tool{
		"gpu_temperature": {
			Name: "gpu_temperature", Description: "Temperature in C of one GPU on one node.", ReadOnly: true,
			Run: func(_ context.Context, args json.RawMessage) (string, error) {
				var a struct {
					Node string `json:"node"`
					GPU  *int   `json:"gpu"`
				}
				if err := decode(args, &a); err != nil {
					return "", fmt.Errorf("bad arguments: %w", err)
				}
				temps, ok := fleet[a.Node]
				if !ok {
					return "", fmt.Errorf("unknown node %q", a.Node)
				}
				if a.GPU == nil || *a.GPU < 0 || *a.GPU >= len(temps) {
					return "", fmt.Errorf("gpu must be between 0 and %d", len(temps)-1)
				}
				return fmt.Sprintf(`{"celsius": %d}`, temps[*a.GPU]), nil
			},
		},
		"drain_node": {
			Name: "drain_node", Description: "Remove a node from service.", ReadOnly: false,
			Run: func(context.Context, json.RawMessage) (string, error) { return `{"drained": true}`, nil },
		},
	}
}

// scriptedModel stands in for an LLM: it looks at the conversation and decides the next move.
func scriptedModel(msgs []Message) Message {
	var results []string
	for _, m := range msgs {
		if m.Role == "tool" {
			results = append(results, m.Content)
		}
	}
	switch len(results) {
	case 0: // first attempt: a wrong argument name, as real models sometimes produce
		return Message{Role: "assistant", Call: &ToolCall{"gpu_temperature", json.RawMessage(`{"node":"n7","index":2}`)}}
	case 1: // it reads the error and corrects itself
		return Message{Role: "assistant", Call: &ToolCall{"gpu_temperature", json.RawMessage(`{"node":"n7","gpu":2}`)}}
	case 2: // it decides to act
		return Message{Role: "assistant", Call: &ToolCall{"drain_node", json.RawMessage(`{"node":"n7"}`)}}
	default:
		return Message{Role: "assistant", Content: "GPU 2 on n7 is at 93 C, above the 85 C limit. Drain result: " + results[len(results)-1]}
	}
}

var ErrStepLimit = errors.New("step limit reached")

// Run is the loop. approve decides whether a tool with side effects may run.
func Run(ctx context.Context, task string, registry map[string]Tool, maxSteps int, approve func(ToolCall) bool) (string, []Message, error) {
	msgs := []Message{
		{Role: "system", Content: "You operate a GPU fleet. Use tools; do not guess."},
		{Role: "user", Content: task},
	}
	for step := 1; step <= maxSteps; step++ {
		if err := ctx.Err(); err != nil {
			return "", msgs, err
		}
		reply := scriptedModel(msgs)
		msgs = append(msgs, reply)
		if reply.Call == nil {
			return reply.Content, msgs, nil // plain text: the task is finished
		}

		call := *reply.Call
		var result string
		tool, known := registry[call.Name]
		switch {
		case !known:
			result = fmt.Sprintf(`{"error": "no tool named %q"}`, call.Name)
		case !tool.ReadOnly && !approve(call):
			result = `{"error": "not approved: this action changes the system"}`
		default:
			tctx, cancel := context.WithTimeout(ctx, 2*time.Second) // a stuck tool cannot hang the task
			out, err := tool.Run(tctx, call.Args)
			cancel()
			if err != nil {
				result = fmt.Sprintf(`{"error": %q}`, err.Error()) // the model sees the error and can retry
			} else {
				result = out
			}
		}
		if len(result) > 2000 {
			result = result[:2000] + "...[truncated]" // never let one result flood the context
		}
		fmt.Printf("  step %d: %s(%s) → %s\n", step, call.Name, call.Args, result)
		msgs = append(msgs, Message{Role: "tool", Content: result})
	}
	return "", msgs, ErrStepLimit
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	task := "Check GPU 2 on node n7 and take it out of service if it is overheating."

	fmt.Println("policy: writes are NOT approved")
	answer, msgs, err := Run(ctx, task, tools(), 6, func(ToolCall) bool { return false })
	fmt.Printf("  answer: %s\n  messages: %d  err: %v\n\n", answer, len(msgs), err)

	fmt.Println("policy: writes approved")
	answer, _, err = Run(ctx, task, tools(), 6, func(c ToolCall) bool { return c.Name == "drain_node" })
	fmt.Printf("  answer: %s\n  err: %v\n\n", answer, err)

	fmt.Println("step limit of 2")
	_, _, err = Run(ctx, task, tools(), 2, func(ToolCall) bool { return true })
	fmt.Println("  err:", err)
}