Automate the Boring Stuff with Go, Ch. 20: Automate with a Local Model

ollama --version
curl http://localhost:11434/api/version

ollama pull qwen2.5:1.5b

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"
)

// Client talks to a model server running on your own machine that speaks the OpenAI
// chat API, which Ollama, llama.cpp's server and LM Studio all do. Nothing leaves the
// machine.
type Client struct {
	BaseURL string // e.g. http://localhost:11434/v1
	Model   string // e.g. qwen2.5:1.5b
	HTTP    *http.Client
}

// FromEnv reads LLM_URL and LLM_MODEL, with Ollama's defaults.
func FromEnv() Client {
	c := Client{BaseURL: "http://localhost:11434/v1", Model: "qwen2.5:1.5b", HTTP: &http.Client{Timeout: 2 * time.Minute}}
	if v := os.Getenv("LLM_URL"); v != "" {
		c.BaseURL = v
	}
	if v := os.Getenv("LLM_MODEL"); v != "" {
		c.Model = v
	}
	return c
}

type message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type responseFormat struct {
	Type string `json:"type"`
}

type chatRequest struct {
	Model          string          `json:"model"`
	Temperature    float64         `json:"temperature"`
	Messages       []message       `json:"messages"`
	ResponseFormat *responseFormat `json:"response_format,omitempty"`
}

type chatResponse struct {
	Choices []struct {
		Message message `json:"message"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

// Chat sends one system prompt and one user prompt and returns the model's text. With
// jsonMode the server is asked for a JSON object; the caller still validates it.
func (c Client) Chat(ctx context.Context, system, user string, jsonMode bool) (string, error) {
	req := chatRequest{Model: c.Model, Temperature: 0, Messages: []message{{"system", system}, {"user", user}}}
	if jsonMode {
		req.ResponseFormat = &responseFormat{Type: "json_object"}
	}
	body, err := json.Marshal(req)
	if err != nil {
		return "", err
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimSuffix(c.BaseURL, "/")+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	resp, err := c.HTTP.Do(httpReq)
	if err != nil {
		return "", fmt.Errorf("model server at %s: %w (is it running? try: ollama serve, then ollama pull %s)", c.BaseURL, err, c.Model)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return "", err
	}
	var out chatResponse
	if err := json.Unmarshal(raw, &out); err != nil || resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("model server: %s: %s", resp.Status, strings.TrimSpace(string(raw)))
	}
	if out.Error != nil {
		return "", errors.New("model server: " + out.Error.Message)
	}
	if len(out.Choices) == 0 {
		return "", errors.New("model server returned no choices")
	}
	return out.Choices[0].Message.Content, nil
}

// ExtractJSON returns the first {...} object in s, for models that wrap their JSON in
// prose or code fences despite being asked not to.
func ExtractJSON(s string) (string, bool) {
	start := strings.Index(s, "{")
	end := strings.LastIndex(s, "}")
	if start < 0 || end <= start {
		return "", false
	}
	return s[start : end+1], true
}

// chat-once sends one question to a model running on your own machine and prints the
// answer. Nothing is sent anywhere else.
//
//	ollama pull qwen2.5:1.5b
//	go run ./ch20/examples/chat-once "In one sentence, what is a goroutine?"
package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"
)

func run(args []string) error {
	if len(args) == 0 {
		return errors.New("usage: chat-once <question>")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	client := FromEnv()
	start := time.Now()
	answer, err := client.Chat(ctx, "You are a concise assistant. Answer in plain text.", strings.Join(args, " "), false)
	if err != nil {
		return err
	}
	fmt.Println(strings.TrimSpace(answer))
	fmt.Fprintf(os.Stderr, "(%s, %s)\n", client.Model, time.Since(start).Round(100*time.Millisecond))
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// structured-output asks the model for JSON in a shape we define, decodes it into a
// struct, and checks every field before believing any of it. When the check fails, the
// model gets one more try with the complaint attached.
//
//	go run ./ch20/examples/structured-output "Picked up 3 kg of apples at the market for 7.50 on Tuesday"
package main

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

// Purchase is the shape we want back. The tags are the contract with the prompt.
type Purchase struct {
	Item     string  `json:"item"`
	Quantity float64 `json:"quantity"`
	Unit     string  `json:"unit"`
	Price    float64 `json:"price"`
}

var units = []string{"kg", "g", "l", "ml", "piece", "pack"}

// Validate is the part the model does not get a vote in.
func (p Purchase) Validate() error {
	var problems []string
	if strings.TrimSpace(p.Item) == "" || len(p.Item) > 60 {
		problems = append(problems, "item must be 1 to 60 characters")
	}
	if p.Quantity <= 0 || p.Quantity > 10000 {
		problems = append(problems, "quantity must be between 0 and 10000")
	}
	if !slices.Contains(units, p.Unit) {
		problems = append(problems, "unit must be one of "+strings.Join(units, ", "))
	}
	if p.Price < 0 || p.Price > 100000 {
		problems = append(problems, "price must be between 0 and 100000")
	}
	if len(problems) > 0 {
		return errors.New(strings.Join(problems, "; "))
	}
	return nil
}

const system = `Extract a purchase from the user's sentence. Reply with JSON only, exactly this shape:
{"item": "<short noun phrase>", "quantity": <number>, "unit": "<one of: kg, g, l, ml, piece, pack>", "price": <number, total paid>}
Use "piece" when no unit is given. Do not add fields or commentary.`

// Extract asks, validates, and retries once with the validation error fed back.
func Extract(ctx context.Context, c Client, sentence string) (Purchase, error) {
	user := sentence
	var last error
	for attempt := 1; attempt <= 2; attempt++ {
		reply, err := c.Chat(ctx, system, user, true)
		if err != nil {
			return Purchase{}, err
		}
		obj, ok := ExtractJSON(reply)
		if !ok {
			last = fmt.Errorf("no JSON object in reply %q", reply)
		} else {
			var p Purchase
			if err := json.Unmarshal([]byte(obj), &p); err != nil {
				last = fmt.Errorf("bad JSON %q: %w", obj, err)
			} else if err := p.Validate(); err != nil {
				last = fmt.Errorf("invalid: %w (reply was %s)", err, obj)
			} else {
				return p, nil
			}
		}
		user = fmt.Sprintf("%s\n\nYour previous answer was rejected: %v. Reply again with valid JSON only.", sentence, last)
	}
	return Purchase{}, fmt.Errorf("gave up after 2 attempts: %w", last)
}

func run(args []string) error {
	if len(args) == 0 {
		return errors.New("usage: structured-output <sentence describing a purchase>")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	p, err := Extract(ctx, FromEnv(), strings.Join(args, " "))
	if err != nil {
		return err
	}
	fmt.Printf("%+v\n", p)
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// categorize-expenses fills in the category column of an expenses CSV by asking a local
// model to pick one from a fixed list. The model proposes; the program decides: an
// answer outside the list becomes "other", every answer can be logged, repeated items are
// asked once, and nothing is written until you have looked at the table.
//
//	go run ./ch20/projects/categorize-expenses ./ch20/projects/categorize-expenses/uncategorized.csv
//	go run ./ch20/projects/categorize-expenses -write ./ch20/projects/categorize-expenses/uncategorized.csv
package main

import (
	"context"
	"encoding/csv"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"os"
	"slices"
	"strings"
	"text/tabwriter"
	"time"
)

// Categories is the allowlist. The model never adds to it.
var Categories = []string{"food", "transport", "books", "household", "health", "entertainment", "other"}

// Row is one expense before and after.
type Row struct {
	Date, Item, Amount, Category string
}

// Prompt is what the model sees for one item: the allowed answers and the exact JSON
// shape, and nothing about the program, the file, or what happens next.
func Prompt(item, amount string) (system, user string) {
	system = fmt.Sprintf(`You classify personal expenses. Reply with JSON only: {"category": "<one of: %s>"}`, strings.Join(Categories, ", "))
	user = fmt.Sprintf("Item: %q, amount: %s", item, amount)
	return system, user
}

// Parse turns the model's reply into a category from the allowlist, or "other" with a
// reason. It never returns anything that is not in Categories.
func Parse(reply string) (category, reason string) {
	obj, ok := ExtractJSON(reply)
	if !ok {
		return "other", "no JSON in reply"
	}
	var v struct {
		Category string `json:"category"`
	}
	if err := json.Unmarshal([]byte(obj), &v); err != nil {
		return "other", "bad JSON"
	}
	c := strings.ToLower(strings.TrimSpace(v.Category))
	if !slices.Contains(Categories, c) {
		return "other", fmt.Sprintf("%q is not an allowed category", c)
	}
	return c, ""
}

// Classifier asks the model, with a memory so the same item is asked once.
type Classifier struct {
	Client Client
	cache  map[string]string
	Calls  int
}

// Category returns the category for one item, asking the model only the first time.
func (k *Classifier) Category(ctx context.Context, item, amount string) string {
	if k.cache == nil {
		k.cache = map[string]string{}
	}
	key := strings.ToLower(strings.TrimSpace(item))
	if c, ok := k.cache[key]; ok {
		return c
	}
	system, user := Prompt(item, amount)
	k.Calls++
	reply, err := k.Client.Chat(ctx, system, user, true)
	if err != nil {
		slog.Warn("model call failed", "item", item, "err", err)
		k.cache[key] = "other"
		return "other"
	}
	c, reason := Parse(reply)
	slog.Info("classified", "item", item, "category", c, "reply", strings.TrimSpace(reply), "note", reason)
	k.cache[key] = c
	return c
}

// ReadRows reads date,item,amount[,category] CSV, columns in any order.
func ReadRows(r io.Reader) ([]Row, error) {
	records, err := csv.NewReader(r).ReadAll()
	if err != nil {
		return nil, err
	}
	if len(records) < 2 {
		return nil, errors.New("no data rows")
	}
	col := map[string]int{}
	for i, h := range records[0] {
		col[strings.ToLower(strings.TrimSpace(h))] = i
	}
	for _, need := range []string{"date", "item", "amount"} {
		if _, ok := col[need]; !ok {
			return nil, fmt.Errorf("missing column %q", need)
		}
	}
	var rows []Row
	for _, rec := range records[1:] {
		row := Row{Date: rec[col["date"]], Item: rec[col["item"]], Amount: rec[col["amount"]]}
		if i, ok := col["category"]; ok && i < len(rec) {
			row.Category = rec[i]
		}
		rows = append(rows, row)
	}
	return rows, nil
}

// Fill classifies every row that has no category yet.
func Fill(ctx context.Context, k *Classifier, rows []Row) {
	for i := range rows {
		if rows[i].Category == "" {
			rows[i].Category = k.Category(ctx, rows[i].Item, rows[i].Amount)
		}
	}
}

func writeRows(w io.Writer, rows []Row) error {
	cw := csv.NewWriter(w)
	cw.Write([]string{"date", "item", "amount", "category"})
	for _, r := range rows {
		cw.Write([]string{r.Date, r.Item, r.Amount, r.Category})
	}
	cw.Flush()
	return cw.Error()
}

func run(args []string) error {
	flags := flag.NewFlagSet("categorize-expenses", flag.ContinueOnError)
	write := flags.Bool("write", false, "write <name>-categorized.csv (default: print the table)")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 1 {
		return errors.New("usage: categorize-expenses [-write] <expenses.csv>")
	}
	f, err := os.Open(flags.Arg(0))
	if err != nil {
		return err
	}
	rows, err := ReadRows(f)
	f.Close()
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	k := &Classifier{Client: FromEnv()}
	start := time.Now()
	Fill(ctx, k, rows)

	tw := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
	fmt.Fprintln(tw, "date\titem\tamount\tcategory")
	for _, r := range rows {
		fmt.Fprintf(tw, "%s\t%s\t%s\t%s\n", r.Date, r.Item, r.Amount, r.Category)
	}
	tw.Flush()
	fmt.Fprintf(os.Stderr, "%d row(s), %d model call(s), %s\n", len(rows), k.Calls, time.Since(start).Round(100*time.Millisecond))

	if *write {
		out := strings.TrimSuffix(flags.Arg(0), ".csv") + "-categorized.csv"
		of, err := os.Create(out)
		if err != nil {
			return err
		}
		if err := writeRows(of, rows); err != nil {
			of.Close()
			return err
		}
		if err := of.Close(); err != nil {
			return err
		}
		fmt.Println("wrote", out)
	}
	return nil
}

func main() {
	level := slog.LevelWarn
	if os.Getenv("LLM_LOG") != "" {
		level = slog.LevelInfo // LLM_LOG=1 shows every prompt's answer
	}
	slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: level})))
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// name-notes suggests a descriptive file name for every text note in a folder by asking
// a local model what each note is about. The model only ever supplies a candidate name;
// a strict pattern decides whether it is acceptable, collisions are resolved by the
// program, and nothing is renamed without -apply.
//
//	go run ./ch20/projects/name-notes ./ch20/projects/name-notes/sample
//	go run ./ch20/projects/name-notes -apply ./ch20/projects/name-notes/sample
package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

// validName is the only shape a new name may take: two to six lowercase words joined by
// single hyphens. No path separators, no dots, no spaces, nothing surprising.
var validName = regexp.MustCompile(`^[a-z0-9]+(-[a-z0-9]+){1,5}$`)

const system = `You name text files. Given the contents of a note, reply with JSON only:
{"name": "<2 to 6 lowercase words joined by hyphens that say what the note is about>"}`

// Suggest asks the model for a name and returns it only if it passes the pattern.
func Suggest(ctx context.Context, c Client, contents string) (string, error) {
	// a note can be long; the first few hundred characters say what it is about
	excerpt := contents
	if len(excerpt) > 600 {
		excerpt = excerpt[:600]
	}
	reply, err := c.Chat(ctx, system, "Note contents:\n"+excerpt, true)
	if err != nil {
		return "", err
	}
	return Accept(reply)
}

// Accept extracts the name from the reply, tidies it the way a person would (lower
// case, hyphens for spaces and underscores, punctuation dropped) and then checks it
// against the pattern. The tidying is deterministic code; the pattern is still the gate.
func Accept(reply string) (string, error) {
	obj, ok := ExtractJSON(reply)
	if !ok {
		return "", fmt.Errorf("no JSON in reply %q", reply)
	}
	var v struct {
		Name string `json:"name"`
	}
	if err := json.Unmarshal([]byte(obj), &v); err != nil {
		return "", fmt.Errorf("bad JSON %q", obj)
	}
	name := Tidy(v.Name)
	if !validName.MatchString(name) {
		return "", fmt.Errorf("suggested name %q does not match %s", v.Name, validName)
	}
	return name, nil
}

// Tidy lower-cases, turns runs of spaces, underscores and hyphens into one hyphen, and
// drops every other character that is not a letter or digit.
func Tidy(s string) string {
	var b strings.Builder
	lastHyphen := true
	for _, r := range strings.ToLower(strings.TrimSpace(s)) {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
			b.WriteRune(r)
			lastHyphen = false
		case r == ' ' || r == '_' || r == '-':
			if !lastHyphen {
				b.WriteByte('-')
				lastHyphen = true
			}
		}
	}
	return strings.TrimSuffix(b.String(), "-")
}

// Rename is one planned change.
type Rename struct {
	From, To string
	Note     string // why there is no To, when there is not
}

// Plan proposes a new name for every .txt and .md file in folder. Names the model gets
// wrong are left as they are, with the reason; duplicate suggestions get a -2, -3 suffix.
func Plan(ctx context.Context, c Client, folder string) ([]Rename, error) {
	entries, err := os.ReadDir(folder)
	if err != nil {
		return nil, err
	}
	taken := map[string]bool{}
	for _, e := range entries {
		taken[e.Name()] = true
	}
	var plan []Rename
	for _, e := range entries {
		ext := strings.ToLower(filepath.Ext(e.Name()))
		if e.IsDir() || (ext != ".txt" && ext != ".md") {
			continue
		}
		contents, err := os.ReadFile(filepath.Join(folder, e.Name()))
		if err != nil {
			return nil, err
		}
		name, err := Suggest(ctx, c, string(contents))
		if err != nil {
			plan = append(plan, Rename{From: e.Name(), Note: err.Error()})
			continue
		}
		candidate := name + ext
		for n := 2; taken[candidate]; n++ {
			candidate = fmt.Sprintf("%s-%d%s", name, n, ext)
		}
		taken[candidate] = true
		plan = append(plan, Rename{From: e.Name(), To: candidate})
	}
	return plan, nil
}

func run(args []string) error {
	flags := flag.NewFlagSet("name-notes", flag.ContinueOnError)
	apply := flags.Bool("apply", false, "rename the files (default: print the plan)")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 1 {
		return errors.New("usage: name-notes [-apply] <folder>")
	}
	folder := flags.Arg(0)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	plan, err := Plan(ctx, FromEnv(), folder)
	if err != nil {
		return err
	}
	for _, r := range plan {
		if r.To == "" {
			fmt.Printf("%-14s kept (%s)\n", r.From, r.Note)
			continue
		}
		fmt.Printf("%-14s -> %s\n", r.From, r.To)
		if *apply {
			if err := os.Rename(filepath.Join(folder, r.From), filepath.Join(folder, r.To)); err != nil {
				return err
			}
		}
	}
	if !*apply {
		fmt.Println("(dry run; pass -apply to rename)")
	}
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}