LLM API Integration Patterns for Backend Engineers

package llm
 
import (
	"context"
	"io"
)
 
type Provider interface {
	Complete(ctx context.Context, req *CompletionRequest) (*CompletionResponse, error)
	Stream(ctx context.Context, req *CompletionRequest) (StreamReader, error)
}
 
type StreamReader interface {
	Next() (StreamChunk, error)
	Close() error
}
 
type CompletionRequest struct {
	// JSON tags are required: without them json.Marshal emits "Model"/"MaxTokens"
	// and the API 400s. These field names (model, messages, max_tokens, stream …)
	// are shared by OpenAI and Anthropic, so the type stays provider-neutral;
	// provider-specific quirks belong in each provider's request mapping.
	Model       string           `json:"model"`
	Messages    []Message        `json:"messages"`
	Tools       []ToolDefinition `json:"tools,omitempty"`
	MaxTokens   int              `json:"max_tokens,omitempty"`
	Temperature float64          `json:"temperature,omitempty"`
	Stream      bool             `json:"stream,omitempty"`
}
 
type CompletionResponse struct {
	Content    string
	ToolCalls  []ToolCall
	Usage      TokenUsage
	Model      string
	FinishReason string
}
 
type TokenUsage struct {
	PromptTokens     int
	CompletionTokens int
	TotalTokens      int
}

type OpenAIProvider struct {
	apiKey     string
	baseURL    string
	httpClient *http.Client
}
 
func NewOpenAIProvider(apiKey string) *OpenAIProvider {
	return &OpenAIProvider{
		apiKey:  apiKey,
		baseURL: "https://api.openai.com/v1",
		httpClient: &http.Client{
			Timeout: 120 * time.Second,
			Transport: &http.Transport{
				MaxIdleConns:        100,
				MaxIdleConnsPerHost: 20,
				IdleConnTimeout:     90 * time.Second,
				ForceAttemptHTTP2:   true, // keeps HTTP/2 if a custom dialer/TLS config is added
			},
		},
	}
}
 
func (p *OpenAIProvider) Complete(ctx context.Context, req *CompletionRequest) (*CompletionResponse, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost,
		p.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("build request: %w", err)
	}
	httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
	httpReq.Header.Set("Content-Type", "application/json")
 
	resp, err := p.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("execute: %w", err)
	}
	defer resp.Body.Close()
 
	if resp.StatusCode != http.StatusOK {
		return nil, parseAPIError(resp)
	}
 
	var result openAIResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("decode response: %w", err)
	}
	return result.toCompletionResponse(), nil
}

func (p *OpenAIProvider) Stream(ctx context.Context, req *CompletionRequest) (StreamReader, error) {
	req.Stream = true
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost,
		p.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("build request: %w", err)
	}
	httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
	httpReq.Header.Set("Accept", "text/event-stream")
 
	resp, err := p.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("execute: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return nil, parseAPIError(resp)
	}
 
	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	return &openAIStreamReader{scanner: scanner, resp: resp}, nil
}
 
func (r *openAIStreamReader) Next() (StreamChunk, error) {
	for r.scanner.Scan() {
		line := r.scanner.Text()
		if line == "" || !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := strings.TrimPrefix(line, "data: ")
		if data == "[DONE]" {
			return StreamChunk{}, io.EOF
		}
		var chunk openAIStreamChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			return StreamChunk{}, fmt.Errorf("decode stream chunk: %w", err)
		}
		return chunk.toStreamChunk(), nil
	}
	// Scan also stops on read errors (connection reset, a line over the 1 MB
	// buffer). Either that or a body that ends before [DONE] is a truncated
	// answer, never a clean io.EOF.
	if err := r.scanner.Err(); err != nil {
		return StreamChunk{}, fmt.Errorf("read stream: %w", err)
	}
	return StreamChunk{}, io.ErrUnexpectedEOF
}

func (h *ChatHandler) StreamCompletion(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")
	w.Header().Set("X-Accel-Buffering", "no") // disable nginx buffering
 
	stream, err := h.provider.Stream(r.Context(), parseRequest(r))
	if err != nil {
		http.Error(w, "upstream failed", http.StatusBadGateway)
		return
	}
	defer stream.Close()
 
	send := func(payload []byte) error {
		if _, err := fmt.Fprintf(w, "data: %s\n\n", payload); err != nil {
			return err
		}
		flusher.Flush()
		return nil
	}
 
	for {
		chunk, err := stream.Next()
		if errors.Is(err, io.EOF) {
			if err := send([]byte("[DONE]")); err != nil {
				h.logger.Info("client disconnected", "err", err)
			}
			return
		}
		if err != nil {
			// No [DONE]: its absence tells the client the answer is truncated.
			h.logger.Warn("upstream stream failed", "err", err)
			return
		}
		data, err := json.Marshal(chunk)
		if err != nil {
			h.logger.Error("marshal chunk", "err", err)
			return
		}
		if err := send(data); err != nil {
			h.logger.Info("client disconnected", "err", err)
			return
		}
	}
}

type RetryProvider struct {
	inner  Provider
	config RetryConfig
	logger *slog.Logger
}
 
type RetryConfig struct {
	MaxRetries     int
	InitialBackoff time.Duration
	MaxBackoff     time.Duration
}
 
func (r *RetryProvider) Complete(ctx context.Context, req *CompletionRequest) (*CompletionResponse, error) {
	var lastErr error
 
	for attempt := 0; attempt <= r.config.MaxRetries; attempt++ {
		if attempt > 0 {
			backoff := r.calculateBackoff(attempt, lastErr)
			r.logger.Info("retrying", "attempt", attempt, "backoff", backoff)
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(backoff):
			}
		}
 
		resp, err := r.inner.Complete(ctx, req)
		if err == nil {
			return resp, nil
		}
 
		lastErr = err
		if !isRetryable(err) {
			return nil, err // fail immediately on non-retryable errors
		}
	}
 
	return nil, fmt.Errorf("exhausted %d retries: %w", r.config.MaxRetries, lastErr)
}
 
func (r *RetryProvider) calculateBackoff(attempt int, err error) time.Duration {
	// Respect Retry-After header
	if apiErr, ok := err.(*APIError); ok && apiErr.RetryAfter > 0 {
		return apiErr.RetryAfter
	}
 
	// Exponential backoff with full jitter
	base := float64(r.config.InitialBackoff) * math.Pow(2, float64(attempt-1))
	if base > float64(r.config.MaxBackoff) {
		base = float64(r.config.MaxBackoff)
	}
	return time.Duration(rand.Float64() * base)
}
 
func isRetryable(err error) bool {
	var apiErr *APIError
	if errors.As(err, &apiErr) {
		switch apiErr.StatusCode {
		case 429:
			// Quota/spend 429s (OpenAI billing error.code values, Anthropic's
			// spend-cap 429 without retry-after) keep failing until access
			// resumes; parseAPIError sets QuotaExhausted for them.
			return !apiErr.QuotaExhausted
		case 500, 502, 503, 504, 529: // 529: Anthropic overloaded_error
			return true
		}
		return false
	}
	return !errors.Is(err, context.Canceled)
}

type FallbackProvider struct {
	chain []ProviderWithModel
}
 
type ProviderWithModel struct {
	Provider Provider
	Model    string
}
 
func (f *FallbackProvider) Complete(ctx context.Context, req *CompletionRequest) (*CompletionResponse, error) {
	var lastErr error
	for _, pm := range f.chain {
		reqCopy := *req
		reqCopy.Model = pm.Model
		resp, err := pm.Provider.Complete(ctx, &reqCopy)
		if err == nil {
			return resp, nil
		}
		lastErr = err
	}
	return nil, fmt.Errorf("all providers failed: %w", lastErr)
}

import "github.com/pkoukk/tiktoken-go"
 
// tiktoken ships OpenAI encodings only; for Claude, use Anthropic's
// token-counting endpoint instead.
type TokenCounter struct {
	mu        sync.Mutex
	encodings map[string]*tiktoken.Tiktoken // building an encoder is the slow part
}
 
func (tc *TokenCounter) encoder(model string) (*tiktoken.Tiktoken, error) {
	tc.mu.Lock()
	defer tc.mu.Unlock()
	if enc, ok := tc.encodings[model]; ok {
		return enc, nil
	}
	enc, err := tiktoken.EncodingForModel(model)
	if err != nil {
		return nil, fmt.Errorf("tokenizer for %s: %w", model, err)
	}
	if tc.encodings == nil {
		tc.encodings = make(map[string]*tiktoken.Tiktoken)
	}
	tc.encodings[model] = enc
	return enc, nil
}
 
func (tc *TokenCounter) CountMessages(model string, messages []Message) (int, error) {
	// gpt-4o-family overheads from OpenAI's "How to count tokens with
	// tiktoken" cookbook: 3 per message, 3 to prime the reply.
	const perMessage = 3
	const replyPriming = 3
 
	enc, err := tc.encoder(model)
	if err != nil {
		return 0, err
	}
	total := replyPriming
	for _, msg := range messages {
		total += perMessage + len(enc.Encode(msg.Role, nil, nil)) + len(enc.Encode(msg.Content, nil, nil))
	}
	return total, nil
}

func SlidingWindow(messages []Message, maxTokens int, counter *TokenCounter, model string) ([]Message, error) {
	var system, conversation []Message
	for _, m := range messages {
		if m.Role == "system" {
			system = append(system, m)
		} else {
			conversation = append(conversation, m)
		}
	}
 
	result := append([]Message(nil), system...)
	budget := maxTokens
	systemTokens, err := counter.CountMessages(model, system)
	if err != nil {
		return nil, fmt.Errorf("count system prompt: %w", err)
	}
	budget -= systemTokens
 
	// Walk newest-to-oldest, keeping each turn that still fits the budget.
	// Prepend kept turns into `kept` so chronological order is preserved;
	// never mutate `conversation` while iterating it.
	var kept []Message
	for i := len(conversation) - 1; i >= 0 && budget > 0; i-- {
		msgTokens, err := counter.CountMessages(model, []Message{conversation[i]})
		if err != nil {
			return nil, fmt.Errorf("count turn %d: %w", i, err)
		}
		if msgTokens > budget {
			break
		}
		kept = append([]Message{conversation[i]}, kept...)
		budget -= msgTokens
	}
	return append(result, kept...), nil
}

var (
	tokenUsage = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "llm_tokens_total",
	}, []string{"model", "type"})
 
	requestCost = promauto.NewCounterVec(prometheus.CounterOpts{
		Name: "llm_cost_dollars_total",
	}, []string{"model"})
)
 
// USD per 1M input/output tokens. Keys must be the exact model IDs you send.
var pricing = map[string]struct{ Input, Output float64 }{
	"gpt-4o":           {2.50, 10.00},
	"gpt-4o-mini":      {0.15, 0.60},
	"claude-sonnet-5":  {2.00, 10.00},
	"claude-haiku-4-5": {1.00, 5.00},
}
 
// killAt trips the switch before 100%: in-flight requests still bill.
const killAt = 0.90
 
type CostTracker struct {
	monthlyBudget float64
	currentSpend  float64
	periodMonth   time.Month // resets currentSpend when the calendar month rolls
	mu            sync.Mutex
	killSwitch    func()
}
 
func (ct *CostTracker) Record(model string, usage TokenUsage) error {
	price, ok := pricing[model]
	if !ok {
		// An unpriced model must not count as $0—that silently disables the kill switch.
		return fmt.Errorf("no price for model %q", model)
	}
	cost := (float64(usage.PromptTokens)/1_000_000)*price.Input +
		(float64(usage.CompletionTokens)/1_000_000)*price.Output
 
	tokenUsage.WithLabelValues(model, "prompt").Add(float64(usage.PromptTokens))
	tokenUsage.WithLabelValues(model, "completion").Add(float64(usage.CompletionTokens))
	requestCost.WithLabelValues(model).Add(cost)
 
	ct.mu.Lock()
	defer ct.mu.Unlock()
 
	// Roll the window each calendar month so "monthly budget" is real — without
	// this, currentSpend only ever grows and the kill switch stays tripped after
	// month one. This counter is per-replica; the authoritative cross-fleet
	// monthly figure is sum(requestCost) in Prometheus, so drive fleet-wide
	// enforcement off that (or a shared store) rather than this local total.
	if m := time.Now().UTC().Month(); m != ct.periodMonth {
		ct.periodMonth = m
		ct.currentSpend = 0
	}
 
	ct.currentSpend += cost
	if ct.currentSpend >= killAt*ct.monthlyBudget && ct.killSwitch != nil {
		slog.Error("LLM budget kill switch tripped", "spend", ct.currentSpend, "budget", ct.monthlyBudget)
		ct.killSwitch()
	}
	return nil
}

type ToolDefinition struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	Parameters  json.RawMessage `json:"parameters"` // JSON Schema
}
 
type ToolCall struct {
	ID        string
	Name      string
	Arguments string
}
 
type Registry struct {
	tools    map[string]ToolDefinition
	schemas  map[string]*jsonschema.Schema
	handlers map[string]func(context.Context, json.RawMessage) (string, error)
}
 
func NewRegistry() *Registry {
	return &Registry{
		tools:    make(map[string]ToolDefinition),
		schemas:  make(map[string]*jsonschema.Schema),
		handlers: make(map[string]func(context.Context, json.RawMessage) (string, error)),
	}
}
 
func (r *Registry) Register(def ToolDefinition, handler func(context.Context, json.RawMessage) (string, error)) error {
	doc, err := jsonschema.UnmarshalJSON(bytes.NewReader(def.Parameters))
	if err != nil {
		return fmt.Errorf("parse schema for %s: %w", def.Name, err)
	}
	// A bare name would resolve to file://<cwd>/<name> and leak server paths
	// into the validation errors the model sees; use a scheme-qualified URL.
	schemaURL := "mem://tools/" + def.Name
	compiler := jsonschema.NewCompiler()
	if err := compiler.AddResource(schemaURL, doc); err != nil {
		return fmt.Errorf("add schema resource for %s: %w", def.Name, err)
	}
	schema, err := compiler.Compile(schemaURL)
	if err != nil {
		return fmt.Errorf("compile schema for %s: %w", def.Name, err)
	}
 
	r.tools[def.Name] = def
	r.schemas[def.Name] = schema
	r.handlers[def.Name] = handler
	return nil
}
 
func (r *Registry) Execute(ctx context.Context, call ToolCall) string {
	handler, ok := r.handlers[call.Name]
	if !ok {
		return "unknown tool"
	}
 
	// Parsing only proves call.Arguments is syntactically valid JSON—it says
	// nothing about whether an order ID is the right type or a value is in
	// range. Extra fields pass too unless the schema sets
	// "additionalProperties": false (JSON Schema allows them by default).
	// Validate against the compiled schema before the handler sees anything.
	doc, err := jsonschema.UnmarshalJSON(strings.NewReader(call.Arguments))
	if err != nil {
		return "invalid arguments: malformed JSON"
	}
	if err := r.schemas[call.Name].Validate(doc); err != nil {
		return "invalid arguments: " + err.Error()
	}
 
	// 30-second timeout on tool execution
	toolCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
 
	result, err := handler(toolCtx, json.RawMessage(call.Arguments))
	if err != nil {
		return "execution failed: " + err.Error()
	}
	return result
}

type backpressuredWriter struct {
	w        http.ResponseWriter
	deadline time.Duration
}
 
func (bw *backpressuredWriter) writeChunk(ctx context.Context, chunk StreamChunk) error {
	data, err := json.Marshal(chunk)
	if err != nil {
		return fmt.Errorf("marshal chunk: %w", err)
	}
 
	// Per-write deadline detects slow clients without blocking forever.
	// SetWriteDeadline is exposed via the http.ResponseController in Go 1.20+.
	rc := http.NewResponseController(bw.w)
	if err := rc.SetWriteDeadline(time.Now().Add(bw.deadline)); err != nil {
		return fmt.Errorf("set deadline: %w", err)
	}
 
	if _, err := fmt.Fprintf(bw.w, "data: %s\n\n", data); err != nil {
		return fmt.Errorf("write: %w", err) // client disconnected or stalled
	}
	// Small chunks sit in the response buffer; the socket write (and the
	// deadline error) happens here. http.Flusher.Flush would drop that error.
	if err := rc.Flush(); err != nil {
		return fmt.Errorf("flush: %w", err)
	}
 
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
		return nil
	}
}

func (h *ChatHandler) StreamWithBackpressure(w http.ResponseWriter, r *http.Request) {
	_, 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")
	w.Header().Set("X-Accel-Buffering", "no")
 
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
 
	stream, err := h.provider.Stream(ctx, parseRequest(r))
	if err != nil {
		http.Error(w, "upstream failed", http.StatusBadGateway)
		return
	}
	defer stream.Close()
 
	bw := &backpressuredWriter{w: w, deadline: 10 * time.Second}
	queue := make(chan StreamChunk, 16) // bounded; fills if client is slow
	var upstreamErr error                // set before close(queue), read after the range ends
 
	go func() {
		defer close(queue)
		for {
			chunk, err := stream.Next()
			if errors.Is(err, io.EOF) {
				return
			}
			if err != nil {
				upstreamErr = err
				return
			}
			select {
			case queue <- chunk:
			case <-ctx.Done():
				return // client disconnected — stop reading from provider
			}
		}
	}()
 
	for chunk := range queue {
		if err := bw.writeChunk(ctx, chunk); err != nil {
			h.logger.Info("client disconnected", "err", err)
			cancel() // unblock the producer goroutine
			return
		}
	}
	if upstreamErr != nil {
		// No [DONE]: its absence tells the client the answer is truncated.
		h.logger.Warn("upstream read failed", "err", upstreamErr)
		return
	}
	if _, err := fmt.Fprintf(w, "data: [DONE]\n\n"); err != nil {
		h.logger.Info("client disconnected", "err", err)
		return
	}
	// Flusher.Flush would drop a stalled write's error here too — see writeChunk.
	if err := http.NewResponseController(w).Flush(); err != nil {
		h.logger.Info("client disconnected", "err", err)
	}
}