Machine Learning & LLMs
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
- No timeout — hung Ollama/vLLM pins goroutines and exhausts the server.
- Loading whole corpora into RAM — use a vector store; page chunk text by id.
- Prompt injection — treat retrieved docs and user text as hostile; separate system instructions clearly; don’t execute tools from doc content.
- Python for the API layer — GIL and packaging pain for a concurrent gateway you could write in Go.
- Logging full prompts with secrets — redaction policies for PII and API keys.
- Unbounded agent loops — cap steps and total tokens per request.
Exercises
- Ollama client — Write
generate(ctx, prompt)against local Ollama; add a 30s timeout test using a cancelled context. - Cosine top-k — Implement
topKwith table-driven tests (identical vectors, orthogonal, empty). - SSE chat — Stream fake tokens to
curl -N; cancel mid-stream and confirm the handler exits. - Tool registry — Register
weatherandtimetools; reject unknown names; unit-test argument JSON validation. - RAG prompt — Given three chunks and a question, build a prompt and assert citations instructions are present.
- Usage metrics — Wrap a completion call and emit slog fields for tokens and latency; scrape logs in a test with
slog.NewJSONHandlerto 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-sseSave 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-toolsSave 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-coreSave 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
tokensanddur_msaroundcomplete.