go-llm/client.go

// Package gollm is a chat LLM wire client: Chat, Stream, and List Models
// through an Ollama-native or OpenAI-compatible Adapter.
package gollm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// Adapter selects the wire protocol a Client speaks.
// [AdapterOpenAI] uses POST /v1/chat/completions and GET /v1/models.
// [AdapterOllama] uses POST /api/chat and GET /api/tags.
type Adapter string

const (
	// AdapterOpenAI selects the OpenAI-compatible protocol.
	AdapterOpenAI Adapter = "openai"
	// AdapterOllama selects the Ollama native protocol.
	AdapterOllama Adapter = "ollama"
)

// Role is the speaker of one Message.
type Role string

const (
	// RoleSystem is a system instruction.
	RoleSystem Role = "system"
	// RoleUser is the caller's turn.
	RoleUser Role = "user"
	// RoleAssistant is a previous model turn.
	RoleAssistant Role = "assistant"
)

// Message is one chat turn. Role and Content are sent unchanged.
type Message struct {
	Role    Role
	Content string
}

// Options configures one Chat or Stream call.
// Model is the server-side model name. The Client does not check that it is
// set; the server rejects an empty model.
// A nil Temperature or MaxTokens is left out of the request.
// On [AdapterOllama], MaxTokens is sent as num_predict.
type Options struct {
	Model       string
	Temperature *float64
	MaxTokens   *int
}

// Client calls one Adapter at one base URL.
// An empty API key sends no Authorization header.
// A non-empty key is sent as a Bearer token on every request.
type Client struct {
	adapter Adapter
	baseURL string
	apiKey  string
	http    *http.Client
}

// NewClient checks adapter and baseURL.
// It returns an error when adapter is not [AdapterOpenAI] or [AdapterOllama],
// or when baseURL is empty or only whitespace.
// A trailing slash on baseURL is removed. apiKey may be empty.
func NewClient(adapter Adapter, baseURL, apiKey string) (*Client, error) {
	switch adapter {
	case AdapterOpenAI, AdapterOllama:
	default:
		return nil, fmt.Errorf("unknown Adapter %q", adapter)
	}
	base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
	if base == "" {
		return nil, fmt.Errorf("base URL is required")
	}
	return &Client{adapter: adapter, baseURL: base, apiKey: apiKey, http: http.DefaultClient}, nil
}

// Chat sends messages and returns the assistant text.
// It returns an error when ctx is canceled, the HTTP status is not 2xx,
// the body cannot be decoded, or an OpenAI response contains no choices.
// An HTTP failure includes a short prefix of the response body.
func (c *Client) Chat(ctx context.Context, messages []Message, opt Options) (string, error) {
	switch c.adapter {
	case AdapterOpenAI:
		return c.openaiChat(ctx, messages, opt)
	case AdapterOllama:
		return c.ollamaChat(ctx, messages, opt)
	default:
		return "", fmt.Errorf("unknown Adapter %q", c.adapter)
	}
}

// Stream sends the same request as Chat with streaming enabled and calls
// yield for each non-empty text chunk, in order.
// It returns the error from yield.
// It also returns an error when ctx is canceled, the HTTP status is not 2xx,
// or a stream line cannot be decoded.
// [AdapterOpenAI] ends the stream when a data line is the token DONE.
// [AdapterOllama] reads one JSON object per line and yields message content.
func (c *Client) Stream(ctx context.Context, messages []Message, opt Options, yield func(chunk string) error) error {
	switch c.adapter {
	case AdapterOpenAI:
		return c.openaiStream(ctx, messages, opt, yield)
	case AdapterOllama:
		return c.ollamaStream(ctx, messages, opt, yield)
	default:
		return fmt.Errorf("unknown Adapter %q", c.adapter)
	}
}

// ListModels returns model names reported by the server.
// [AdapterOpenAI] reads data[].id from GET /v1/models.
// [AdapterOllama] reads models[].name from GET /api/tags.
// It returns an error when ctx is canceled, the HTTP status is not 2xx,
// or the body cannot be decoded. Empty names are omitted.
func (c *Client) ListModels(ctx context.Context) ([]string, error) {
	switch c.adapter {
	case AdapterOpenAI:
		return c.openaiListModels(ctx)
	case AdapterOllama:
		return c.ollamaListModels(ctx)
	default:
		return nil, fmt.Errorf("unknown Adapter %q", c.adapter)
	}
}

func (c *Client) auth(req *http.Request) {
	if c.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
	}
}

type openaiChatRequest struct {
	Model       string          `json:"model"`
	Messages    []openaiMessage `json:"messages"`
	Stream      bool            `json:"stream,omitempty"`
	Temperature *float64        `json:"temperature,omitempty"`
	MaxTokens   *int            `json:"max_tokens,omitempty"`
}

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

func openaiMessages(messages []Message) []openaiMessage {
	out := make([]openaiMessage, len(messages))
	for i, m := range messages {
		out[i] = openaiMessage{Role: string(m.Role), Content: m.Content}
	}
	return out
}

func (c *Client) openaiChat(ctx context.Context, messages []Message, opt Options) (string, error) {
	body, err := json.Marshal(openaiChatRequest{
		Model:       opt.Model,
		Messages:    openaiMessages(messages),
		Temperature: opt.Temperature,
		MaxTokens:   opt.MaxTokens,
	})
	if err != nil {
		return "", err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	c.auth(req)
	res, err := c.http.Do(req)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()
	raw, err := io.ReadAll(res.Body)
	if err != nil {
		return "", err
	}
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		return "", fmt.Errorf("openai chat: %s: %s", res.Status, truncate(raw))
	}
	var parsed struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return "", err
	}
	if len(parsed.Choices) == 0 {
		return "", fmt.Errorf("openai chat: no choices")
	}
	return parsed.Choices[0].Message.Content, nil
}

func (c *Client) openaiStream(ctx context.Context, messages []Message, opt Options, yield func(string) error) error {
	body, err := json.Marshal(openaiChatRequest{
		Model:       opt.Model,
		Messages:    openaiMessages(messages),
		Stream:      true,
		Temperature: opt.Temperature,
		MaxTokens:   opt.MaxTokens,
	})
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream")
	c.auth(req)
	res, err := c.http.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		raw, _ := io.ReadAll(res.Body)
		return fmt.Errorf("openai stream: %s: %s", res.Status, truncate(raw))
	}
	sc := bufio.NewScanner(res.Body)
	sc.Buffer(nil, 1024*1024)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" || !strings.HasPrefix(line, "data:") {
			continue
		}
		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "[DONE]" {
			return nil
		}
		var parsed struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
		}
		if err := json.Unmarshal([]byte(payload), &parsed); err != nil {
			return err
		}
		if len(parsed.Choices) == 0 {
			continue
		}
		chunk := parsed.Choices[0].Delta.Content
		if chunk == "" {
			continue
		}
		if err := yield(chunk); err != nil {
			return err
		}
	}
	return sc.Err()
}

func (c *Client) openaiListModels(ctx context.Context) ([]string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/v1/models", nil)
	if err != nil {
		return nil, err
	}
	c.auth(req)
	res, err := c.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	raw, err := io.ReadAll(res.Body)
	if err != nil {
		return nil, err
	}
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		return nil, fmt.Errorf("openai models: %s: %s", res.Status, truncate(raw))
	}
	var parsed struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, err
	}
	names := make([]string, 0, len(parsed.Data))
	for _, d := range parsed.Data {
		if d.ID != "" {
			names = append(names, d.ID)
		}
	}
	return names, nil
}

func truncate(b []byte) string {
	s := strings.TrimSpace(string(b))
	if len(s) > 200 {
		return s[:200]
	}
	return s
}