Machine Learning & LLMs

Updated

July 30, 2026

AI & LLMs in Go

Overview

Python owns model training. Go owns serving, orchestration, and glue: APIs, RAG pipelines, agents, tool routers, WebSocket streams, and high-concurrency gateways in front of model runtimes.

In 2026 you rarely train foundation models yourself. You call hosted APIs or local runtimes (Ollama, vLLM, llama.cpp servers) and build reliable product software around them. Go’s strengths—static binaries, context cancellation, excellent HTTP—map directly onto that job.

Architecture Split

[Python / CUDA land]          [Go land]
 research notebooks            API gateway
 training / fine-tunes         RAG orchestration
 experimental kernels          auth, quotas, billing
                               tool/function calling
                               workers + queues

Keep models and experiments where the ecosystem is; put the product boundary in Go.

Calling Local Models (Ollama)

Ollama exposes an HTTP API. Many teams use LangChainGo (github.com/tmc/langchaingo) or a thin custom client.

package main

import (
    "context"
    "fmt"
    "log"
    "time"

    "github.com/tmc/langchaingo/llms"
    "github.com/tmc/langchaingo/llms/ollama"
)

func main() {
    llm, err := ollama.New(ollama.WithModel("llama3.2"))
    if err != nil {
        log.Fatal(err)
    }

    ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
    defer cancel()

    out, err := llms.GenerateFromSinglePrompt(ctx, llm, "Explain Go channels in 3 bullets")
    if err != nil {
        log.Fatal(err)
    }
    fmt.Println(out)
}

Minimal raw HTTP client

Prefer owning the client when you need strict control:

type genRequest struct {
    Model  string `json:"model"`
    Prompt string `json:"prompt"`
    Stream bool   `json:"stream"`
}

type genResponse struct {
    Response string `json:"response"`
    Done     bool   `json:"done"`
}

func generate(ctx context.Context, client *http.Client, base, model, prompt string) (string, error) {
    body, _ := json.Marshal(genRequest{Model: model, Prompt: prompt, Stream: false})
    req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/api/generate", bytes.NewReader(body))
    if err != nil {
        return "", err
    }
    req.Header.Set("Content-Type", "application/json")

    res, err := client.Do(req)
    if err != nil {
        return "", err
    }
    defer res.Body.Close()
    if res.StatusCode >= 300 {
        b, _ := io.ReadAll(res.Body)
        return "", fmt.Errorf("ollama %s: %s", res.Status, b)
    }
    var gr genResponse
    if err := json.NewDecoder(res.Body).Decode(&gr); err != nil {
        return "", err
    }
    return gr.Response, nil
}

Always set timeouts and honor ctx cancellation—model calls are slow and easy to pile up under load.

Streaming Tokens to Clients

UX for chat means streaming. Bridge model SSE/NDJSON to your API:

func streamHandler(w http.ResponseWriter, r *http.Request) {
    flusher, ok := w.(http.Flusher)
    if !ok {
        http.Error(w, "streaming unsupported", http.StatusInternalServerError)
        return
    }
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")

    // pseudo: for each token from model stream
    tokens := []string{"Hello", " ", "from", " ", "Go"}
    for _, t := range tokens {
        select {
        case <-r.Context().Done():
            return
        default:
        }
        fmt.Fprintf(w, "data: %s\n\n", t)
        flusher.Flush()
    }
    fmt.Fprintf(w, "data: [DONE]\n\n")
    flusher.Flush()
}

RAG Pipeline in Go

query
  --> embed(query) -> []float32
  --> vector search top-k chunks
  --> build prompt (system + chunks + query)
  --> LLM completion
  --> cite sources + log token usage

Embedding + top-k stub

type Chunk struct {
    ID   string
    Text string
    // Vector stored in DB; not always loaded
}

type Scored struct {
    Chunk
    Score float32
}

func cosine(a, b []float32) float32 {
    var dot, na, nb float32
    for i := range a {
        dot += a[i] * b[i]
        na += a[i] * a[i]
        nb += b[i] * b[i]
    }
    if na == 0 || nb == 0 {
        return 0
    }
    return dot / (float32(math.Sqrt(float64(na))) * float32(math.Sqrt(float64(nb))))
}

func topK(query []float32, docs []struct {
    Chunk
    Vec []float32
}, k int) []Scored {
    scored := make([]Scored, 0, len(docs))
    for _, d := range docs {
        scored = append(scored, Scored{Chunk: d.Chunk, Score: cosine(query, d.Vec)})
    }
    sort.Slice(scored, func(i, j int) bool { return scored[i].Score > scored[j].Score })
    if k > len(scored) {
        k = len(scored)
    }
    return scored[:k]
}

Production retrieval usually lives in pgvector, Weaviate, Milvus, Qdrant—not in-process loops. Go is the language of several of those systems’ clients and control planes.

Prompt assembly

func buildPrompt(system string, chunks []Scored, question string) string {
    var b strings.Builder
    b.WriteString(system)
    b.WriteString("\n\nContext:\n")
    for i, c := range chunks {
        fmt.Fprintf(&b, "[%d] (%s) %s\n", i+1, c.ID, c.Text)
    }
    b.WriteString("\nQuestion: ")
    b.WriteString(question)
    b.WriteString("\nAnswer using only the context. Cite chunk ids.\n")
    return b.String()
}

Tool / Function Calling Agents

Agents are orchestration: model proposes a tool call → Go executes allow-listed tools → model continues.

type Tool interface {
    Name() string
    Description() string
    Run(ctx context.Context, args json.RawMessage) (string, error)
}

type Registry map[string]Tool

func (r Registry) Exec(ctx context.Context, name string, args json.RawMessage) (string, error) {
    t, ok := r[name]
    if !ok {
        return "", fmt.Errorf("unknown tool %q", name)
    }
    // enforce timeout per tool
    ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
    defer cancel()
    return t.Run(ctx, args)
}

Never let the model choose arbitrary shell commands or unconstrained HTTP. Map names to vetted implementations only.

SIMD and Local Vector Math (Experimental)

Go continues to invest in performance for numeric loops (SIMD experiments via GOEXPERIMENT in recent releases). For production similarity at scale you still want a vector DB; for small in-process indexes or filters, careful float32 loops (and assembly/SIMD where justified) help.

# When available in your Go version / experiment flags:
GOEXPERIMENT=simd go test ./internal/vec/...

Profile before micro-optimizing: retrieval latency is usually network + model, not dot products.

Reliability Patterns

Concern Go approach
Slow models context.WithTimeout, client-side deadlines
Overload worker pool / queue; reject with 429
Cost token counters, per-tenant budgets
Safety output filters, tool allow-lists, PII redaction logs
Observability slog attrs: request_id, model, tokens_in/out, latency_ms
func withUsageLog(model string, tokensIn, tokensOut int, d time.Duration) {
    slog.Info("llm.completion",
        "model", model,
        "tokens_in", tokensIn,
        "tokens_out", tokensOut,
        "latency_ms", d.Milliseconds(),
    )
}

Hosted APIs

Same patterns against OpenAI-compatible endpoints:

req, _ := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload))
req.Header.Set("Authorization", "Bearer "+os.Getenv("OPENAI_API_KEY"))
req.Header.Set("Content-Type", "application/json")
  • Store keys in env/secret manager, never in the repo.
  • Prefer one outbound client with connection pooling and explicit timeouts.
  • Implement retries only on idempotent GETs or with careful dedupe keys for jobs—not blindly on every chat POST.

Production Checklist

  • Timeouts on every model HTTP call
  • Cancellation wired from client disconnect to upstream
  • Streaming path tested under proxy buffering settings
  • Tool calls allow-listed and argument-validated
  • RAG sources logged for audit (chunk ids)
  • Token/cost metrics per tenant
  • Secrets only in env/secret store
  • Rate limits at the edge
  • Eval suite for prompt regressions (golden questions)

Common Pitfalls

  1. No timeout — hung Ollama/vLLM pins goroutines and exhausts the server.
  2. Loading whole corpora into RAM — use a vector store; page chunk text by id.
  3. Prompt injection — treat retrieved docs and user text as hostile; separate system instructions clearly; don’t execute tools from doc content.
  4. Python for the API layer — GIL and packaging pain for a concurrent gateway you could write in Go.
  5. Logging full prompts with secrets — redaction policies for PII and API keys.
  6. Unbounded agent loops — cap steps and total tokens per request.

Exercises

  1. Ollama client — Write generate(ctx, prompt) against local Ollama; add a 30s timeout test using a cancelled context.
  2. Cosine top-k — Implement topK with table-driven tests (identical vectors, orthogonal, empty).
  3. SSE chat — Stream fake tokens to curl -N; cancel mid-stream and confirm the handler exits.
  4. Tool registry — Register weather and time tools; reject unknown names; unit-test argument JSON validation.
  5. RAG prompt — Given three chunks and a question, build a prompt and assert citations instructions are present.
  6. Usage metrics — Wrap a completion call and emit slog fields for tokens and latency; scrape logs in a test with slog.NewJSONHandler to a buffer.

More examples

SSE token stream with cancel

mkdir -p /tmp/go-llm-sse && cd /tmp/go-llm-sse
go mod init example.com/llm-sse

Save as main.go:

package main

import (
    "bufio"
    "context"
    "fmt"
    "net/http"
    "net/http/httptest"
    "strings"
    "time"
)

func main() {
    tokens := []string{"Hello", " ", "world", "!"}
    h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        fl, ok := w.(http.Flusher)
        if !ok {
            http.Error(w, "no flush", 500)
            return
        }
        w.Header().Set("Content-Type", "text/event-stream")
        for _, t := range tokens {
            select {
            case <-r.Context().Done():
                return
            default:
            }
            fmt.Fprintf(w, "data: %s\n\n", t)
            fl.Flush()
            time.Sleep(5 * time.Millisecond)
        }
        fmt.Fprintf(w, "data: [DONE]\n\n")
    })

    ts := httptest.NewServer(h)
    defer ts.Close()

    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    req, _ := http.NewRequestWithContext(ctx, http.MethodGet, ts.URL, nil)
    resp, err := http.DefaultClient.Do(req)
    if err != nil {
        panic(err)
    }
    defer resp.Body.Close()

    var got []string
    sc := bufio.NewScanner(resp.Body)
    for sc.Scan() {
        line := sc.Text()
        if strings.HasPrefix(line, "data: ") {
            got = append(got, strings.TrimPrefix(line, "data: "))
        }
    }
    fmt.Println("tokens:", strings.Join(got, "|"))
}
go run .

Expected output:

tokens: Hello| |world|!|[DONE]

Tool registry with JSON args

mkdir -p /tmp/go-llm-tools && cd /tmp/go-llm-tools
go mod init example.com/llm-tools

Save as main.go:

package main

import (
    "encoding/json"
    "fmt"
)

type ToolFunc func(args json.RawMessage) (string, error)

type Registry map[string]ToolFunc

func (r Registry) Call(name string, args json.RawMessage) (string, error) {
    fn, ok := r[name]
    if !ok {
        return "", fmt.Errorf("unknown tool %q", name)
    }
    return fn(args)
}

func main() {
    reg := Registry{
        "echo": func(args json.RawMessage) (string, error) {
            var in struct {
                Text string `json:"text"`
            }
            if err := json.Unmarshal(args, &in); err != nil {
                return "", err
            }
            return in.Text, nil
        },
    }
    out, err := reg.Call("echo", json.RawMessage(`{"text":"hi"}`))
    fmt.Println("echo:", out, err)
    _, err = reg.Call("weather", json.RawMessage(`{}`))
    fmt.Println("missing:", err != nil)
}
go run .

Expected output:

echo: hi <nil>
missing: true

Runnable example

LLM gateways need timeouts, tool registries, and retrieval helpers—not only HTTP clients to vendors. This program implements cosine top-k, a tool allowlist, and a timeout-bounded “completion” (stdlib mock).

mkdir -p /tmp/go-llm-core && cd /tmp/go-llm-core
go mod init example.com/llm-core

Save as main.go:

package main

import (
    "context"
    "encoding/json"
    "fmt"
    "math"
    "sort"
    "strings"
    "time"
)

type Doc struct {
    ID   string
    Text string
    Vec  []float64
}

func cosine(a, b []float64) float64 {
    var dot, na, nb float64
    for i := range a {
        dot += a[i] * b[i]
        na += a[i] * a[i]
        nb += b[i] * b[i]
    }
    if na == 0 || nb == 0 {
        return 0
    }
    return dot / (math.Sqrt(na) * math.Sqrt(nb))
}

func topK(query []float64, docs []Doc, k int) []Doc {
    type scored struct {
        d Doc
        s float64
    }
    var all []scored
    for _, d := range docs {
        all = append(all, scored{d, cosine(query, d.Vec)})
    }
    sort.Slice(all, func(i, j int) bool { return all[i].s > all[j].s })
    if k > len(all) {
        k = len(all)
    }
    out := make([]Doc, 0, k)
    for i := 0; i < k; i++ {
        out = append(out, all[i].d)
    }
    return out
}

type ToolFunc func(ctx context.Context, args json.RawMessage) (string, error)

type Registry struct {
    tools map[string]ToolFunc
}

func NewRegistry() *Registry {
    return &Registry{tools: map[string]ToolFunc{}}
}

func (r *Registry) Register(name string, fn ToolFunc) { r.tools[name] = fn }

func (r *Registry) Call(ctx context.Context, name string, args json.RawMessage) (string, error) {
    fn, ok := r.tools[name]
    if !ok {
        return "", fmt.Errorf("unknown tool %q", name)
    }
    return fn(ctx, args)
}

func complete(ctx context.Context, prompt string) (string, error) {
    // Stand-in for Ollama/OpenAI: honor context deadline.
    select {
    case <-ctx.Done():
        return "", ctx.Err()
    case <-time.After(20 * time.Millisecond):
        return "answer based on: " + truncate(prompt, 60), nil
    }
}

func truncate(s string, n int) string {
    if len(s) <= n {
        return s
    }
    return s[:n] + "..."
}

func main() {
    docs := []Doc{
        {ID: "1", Text: "Go concurrency uses goroutines", Vec: []float64{1, 0, 0}},
        {ID: "2", Text: "Rust has ownership and borrowing", Vec: []float64{0, 1, 0}},
        {ID: "3", Text: "Go interfaces are satisfied implicitly", Vec: []float64{0.9, 0.1, 0}},
    }
    q := []float64{1, 0, 0}
    hits := topK(q, docs, 2)
    fmt.Println("top-k:")
    for _, h := range hits {
        fmt.Printf("  %s: %s\n", h.ID, h.Text)
    }

    var b strings.Builder
    b.WriteString("Use only these sources:\n")
    for _, h := range hits {
        fmt.Fprintf(&b, "- [%s] %s\n", h.ID, h.Text)
    }
    b.WriteString("Question: What does Go use for concurrency?\n")
    prompt := b.String()

    reg := NewRegistry()
    reg.Register("time_now", func(ctx context.Context, _ json.RawMessage) (string, error) {
        return time.Now().UTC().Format(time.RFC3339), nil
    })
    out, err := reg.Call(context.Background(), "time_now", nil)
    fmt.Println("tool time_now:", out, err)
    _, err = reg.Call(context.Background(), "shell_exec", nil)
    fmt.Println("unknown tool rejected:", err != nil)

    ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
    defer cancel()
    ans, err := complete(ctx, prompt)
    fmt.Println("completion:", ans, err)

    // Timeout path
    ctx2, cancel2 := context.WithTimeout(context.Background(), time.Nanosecond)
    defer cancel2()
    time.Sleep(time.Millisecond)
    _, err = complete(ctx2, prompt)
    fmt.Println("timeout err:", err != nil)
}
go run .

Expected output (timestamps vary):

top-k:
  1: Go concurrency uses goroutines
  3: Go interfaces are satisfied implicitly
tool time_now: 2026-07-28T12:00:00Z <nil>
unknown tool rejected: true
completion: answer based on: Use only these sources:
- [1] Go concurrency uses... <nil>
timeout err: true

What to notice

  • Retrieval (top-k) and tool allowlists are pure Go; treat model I/O as an untrusted boundary with timeouts.
  • Unknown tools must fail closed—never reflect arbitrary names into os/exec.
  • Cancelled contexts prevent hung upstream model servers from pinning handlers forever.

Try next

  • Stream fake tokens to stdout with flushes (SSE-style) and cancel mid-stream.
  • Add slog fields for tokens and dur_ms around complete.