go-llm/ollama.go

package gollm

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

type ollamaChatRequest struct {
	Model    string          `json:"model"`
	Messages []openaiMessage `json:"messages"`
	Stream   bool            `json:"stream"`
	Options  map[string]any  `json:"options,omitempty"`
}

func ollamaOptions(opt Options) map[string]any {
	if opt.Temperature == nil && opt.MaxTokens == nil {
		return nil
	}
	m := map[string]any{}
	if opt.Temperature != nil {
		m["temperature"] = *opt.Temperature
	}
	if opt.MaxTokens != nil {
		m["num_predict"] = *opt.MaxTokens
	}
	return m
}

func (c *Client) ollamaChat(ctx context.Context, messages []Message, opt Options) (string, error) {
	body, err := json.Marshal(ollamaChatRequest{
		Model:    opt.Model,
		Messages: openaiMessages(messages),
		Stream:   false,
		Options:  ollamaOptions(opt),
	})
	if err != nil {
		return "", err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/chat", 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("ollama chat: %s: %s", res.Status, truncate(raw))
	}
	var parsed struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	}
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return "", err
	}
	return parsed.Message.Content, nil
}

func (c *Client) ollamaStream(ctx context.Context, messages []Message, opt Options, yield func(string) error) error {
	body, err := json.Marshal(ollamaChatRequest{
		Model:    opt.Model,
		Messages: openaiMessages(messages),
		Stream:   true,
		Options:  ollamaOptions(opt),
	})
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/chat", 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()
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		raw, _ := io.ReadAll(res.Body)
		return fmt.Errorf("ollama 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 == "" {
			continue
		}
		var parsed struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		}
		if err := json.Unmarshal([]byte(line), &parsed); err != nil {
			return err
		}
		if parsed.Message.Content == "" {
			continue
		}
		if err := yield(parsed.Message.Content); err != nil {
			return err
		}
	}
	return sc.Err()
}

func (c *Client) ollamaListModels(ctx context.Context) ([]string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/api/tags", 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("ollama tags: %s: %s", res.Status, truncate(raw))
	}
	var parsed struct {
		Models []struct {
			Name string `json:"name"`
		} `json:"models"`
	}
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, err
	}
	names := make([]string, 0, len(parsed.Models))
	for _, m := range parsed.Models {
		if m.Name != "" {
			names = append(names, m.Name)
		}
	}
	return names, nil
}