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 is which native protocol the Client speaks.
type Adapter string

const (
	AdapterOpenAI Adapter = "openai"
	AdapterOllama Adapter = "ollama"
)

// Role is a Message role.
type Role string

const (
	RoleSystem    Role = "system"
	RoleUser      Role = "user"
	RoleAssistant Role = "assistant"
)

// Message is one chat turn.
type Message struct {
	Role    Role
	Content string
}

// Options are Chat/Stream parameters. Temperature and MaxTokens are omitted when nil.
type Options struct {
	Model       string
	Temperature *float64
	MaxTokens   *int
}

// Client is a configured caller bound to one Adapter, base URL, and optional API key.
type Client struct {
	adapter Adapter
	baseURL string
	apiKey  string
	http    *http.Client
}

// NewClient validates Adapter and base URL. 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 returns the assistant text for messages.
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 yields assistant text chunks for the same Chat request.
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 server-side Model names.
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
}