Calling LLMs

{
  "model": "some-model",
  "messages": [
    {"role": "system", "content": "You are concise."},
    {"role": "user", "content": "Summarize this ticket."}
  ],
  "max_tokens": 400,
  "temperature": 0.2,
  "stream": true
}

// llmclient.go — a streaming chat client: retries with backoff, SSE parsing, structured output.
package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync/atomic"
	"time"
)

// ---- the wire format shared by OpenAI-compatible servers (vLLM, Ollama, llama.cpp, hosted APIs)

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

type ChatRequest struct {
	Model          string          `json:"model"`
	Messages       []Message       `json:"messages"`
	Stream         bool            `json:"stream,omitempty"`
	MaxTokens      int             `json:"max_tokens,omitempty"`
	Temperature    float64         `json:"temperature"`
	ResponseFormat json.RawMessage `json:"response_format,omitempty"`
}

type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
}

type chunk struct {
	Choices []struct {
		Delta        struct{ Content string } `json:"delta"`
		Message      Message                  `json:"message"`
		FinishReason string                   `json:"finish_reason"`
	} `json:"choices"`
	Usage *Usage `json:"usage"`
}

type APIError struct {
	Status     int
	RetryAfter time.Duration
	Body       string
}

func (e *APIError) Error() string { return fmt.Sprintf("api status %d: %s", e.Status, e.Body) }

func (e *APIError) retryable() bool { return e.Status == 429 || e.Status >= 500 }

// ---- client

type Client struct {
	BaseURL, APIKey string
	HTTP            *http.Client // shared: it owns the connection pool
	MaxRetries      int
}

// do sends one request, retrying retryable failures with exponential backoff and jitter.
func (c *Client) do(ctx context.Context, req ChatRequest) (*http.Response, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}
	backoff := 20 * time.Millisecond
	for attempt := 0; ; attempt++ {
		hr, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"/v1/chat/completions", bytes.NewReader(body))
		if err != nil {
			return nil, err
		}
		hr.Header.Set("Content-Type", "application/json")
		hr.Header.Set("Authorization", "Bearer "+c.APIKey)
		resp, err := c.HTTP.Do(hr)
		if err == nil && resp.StatusCode == http.StatusOK {
			return resp, nil
		}
		var apiErr *APIError
		if err == nil {
			msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
			resp.Body.Close() // always: otherwise the connection cannot be reused
			apiErr = &APIError{Status: resp.StatusCode, Body: strings.TrimSpace(string(msg))}
			if s, convErr := strconv.Atoi(resp.Header.Get("Retry-After")); convErr == nil {
				apiErr.RetryAfter = time.Duration(s) * time.Second
			}
			err = apiErr
		}
		if ctx.Err() != nil || attempt >= c.MaxRetries || (apiErr != nil && !apiErr.retryable()) {
			return nil, fmt.Errorf("after %d attempt(s): %w", attempt+1, err)
		}
		wait := backoff + time.Duration(rand.Int63n(int64(backoff))) // jitter: do not retry in step with everyone else
		if apiErr != nil && apiErr.RetryAfter > 0 {
			wait = apiErr.RetryAfter
		}
		fmt.Printf("  attempt %d failed (%v); retrying in %v\n", attempt+1, err, wait.Round(time.Millisecond))
		select {
		case <-time.After(wait):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
		backoff *= 2
	}
}

// Stream calls onToken for each piece of text as it arrives and returns the usage.
func (c *Client) Stream(ctx context.Context, req ChatRequest, onToken func(string)) (Usage, string, error) {
	req.Stream = true
	resp, err := c.do(ctx, req)
	if err != nil {
		return Usage{}, "", err
	}
	defer resp.Body.Close()

	var usage Usage
	finish := ""
	sc := bufio.NewScanner(resp.Body)
	sc.Buffer(make([]byte, 0, 64<<10), 1<<20) // one event can exceed the 64 kB default
	for sc.Scan() {
		data, ok := strings.CutPrefix(sc.Text(), "data: ")
		if !ok {
			continue // blank separators and comments
		}
		if data == "[DONE]" {
			return usage, finish, nil
		}
		var ch chunk
		if err := json.Unmarshal([]byte(data), &ch); err != nil {
			return usage, finish, fmt.Errorf("bad event %q: %w", data, err)
		}
		if ch.Usage != nil {
			usage = *ch.Usage
		}
		for _, choice := range ch.Choices {
			if choice.Delta.Content != "" {
				onToken(choice.Delta.Content)
			}
			if choice.FinishReason != "" {
				finish = choice.FinishReason
			}
		}
	}
	if err := sc.Err(); err != nil {
		return usage, finish, err
	}
	return usage, finish, errors.New("stream ended without [DONE]")
}

// Structured asks for JSON matching a schema and decodes it strictly into out.
func (c *Client) Structured(ctx context.Context, req ChatRequest, schema string, out any) error {
	req.ResponseFormat = json.RawMessage(`{"type":"json_schema","json_schema":{"name":"result","strict":true,"schema":` + schema + `}}`)
	resp, err := c.do(ctx, req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	var ch chunk
	if err := json.NewDecoder(resp.Body).Decode(&ch); err != nil {
		return err
	}
	if len(ch.Choices) == 0 {
		return errors.New("no choices in response")
	}
	dec := json.NewDecoder(strings.NewReader(ch.Choices[0].Message.Content))
	dec.DisallowUnknownFields() // the model's output is untrusted input: validate it
	return dec.Decode(out)
}

// ---- a fake OpenAI-compatible server, so the program runs offline

func fakeServer() *httptest.Server {
	var calls atomic.Int32
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req ChatRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, "bad json", http.StatusBadRequest)
			return
		}
		if len(req.Messages) == 0 {
			http.Error(w, "messages is required", http.StatusBadRequest)
			return
		}
		if calls.Add(1) == 1 { // the first call is rate limited
			http.Error(w, "rate limit exceeded", http.StatusTooManyRequests)
			return
		}
		if !req.Stream {
			content := `{"sentiment":"positive","confidence":0.93}`
			json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{
				"message": Message{"assistant", content}, "finish_reason": "stop"}}})
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		rc := http.NewResponseController(w)
		words := strings.Fields("Go handles the streaming and the retries ; the model does the thinking .")
		for _, word := range words {
			time.Sleep(8 * time.Millisecond)
			ev, _ := json.Marshal(map[string]any{"choices": []any{map[string]any{"delta": map[string]string{"content": word + " "}}}})
			fmt.Fprintf(w, "data: %s\n\n", ev)
			rc.Flush()
		}
		fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":21,\"completion_tokens\":%d}}\n\n", len(words))
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
}

func main() {
	srv := fakeServer()
	defer srv.Close()
	c := &Client{BaseURL: srv.URL, APIKey: "test", HTTP: srv.Client(), MaxRetries: 3}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) // total budget for the call
	defer cancel()

	fmt.Println("streaming call:")
	start := time.Now()
	var first time.Duration
	fmt.Print("  ")
	usage, finish, err := c.Stream(ctx, ChatRequest{
		Model:    "local-model",
		Messages: []Message{{"system", "You are concise."}, {"user", "Who does what?"}},
	}, func(tok string) {
		if first == 0 {
			first = time.Since(start)
			fmt.Print("  ")
		}
		fmt.Print(tok)
	})
	fmt.Printf("\n  finish=%s  prompt=%d  completion=%d tokens  ttft=%dms  total=%dms  err=%v\n",
		finish, usage.PromptTokens, usage.CompletionTokens, first.Milliseconds(), time.Since(start).Milliseconds(), err)

	fmt.Println("structured call:")
	var result struct {
		Sentiment  string  `json:"sentiment"`
		Confidence float64 `json:"confidence"`
	}
	schema := `{"type":"object","properties":{"sentiment":{"type":"string","enum":["positive","negative","neutral"]},"confidence":{"type":"number"}},"required":["sentiment","confidence"],"additionalProperties":false}`
	err = c.Structured(ctx, ChatRequest{Model: "local-model", Messages: []Message{{"user", "Classify: 'this is great'"}}}, schema, &result)
	fmt.Printf("  decoded: %+v  err=%v\n", result, err)

	// A client error is not retried: sending the same bad request again cannot help.
	_, _, err = c.Stream(ctx, ChatRequest{Model: "local-model"}, func(string) {})
	fmt.Println("bad request:", err)
}