petal/internal/llm/ollama.go

package llm

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

// OllamaClient talks to Ollama's native /api/chat endpoint. Selected when
// LLM_BACKEND=ollama.
type OllamaClient struct {
	backend
}

// ollamaRequest mirrors Ollama's /api/chat body. Sampling parameters live under
// "options" (see the cheatsheet in the spec).
//
// Think is sent false unconditionally: reasoning models (e.g. qwen3.5) otherwise
// stream their chain-of-thought into a separate "thinking" field and can exhaust
// num_predict before emitting any answer in "content" — leaving us an empty
// response. We want direct output (structured JSON for the checkpoint, concise
// prose for chat), not reasoning. Non-thinking models simply ignore the flag.
type ollamaRequest struct {
	Model    string        `json:"model"`
	Messages []Message     `json:"messages"`
	Stream   bool          `json:"stream"`
	Think    bool          `json:"think"`
	Options  ollamaOptions `json:"options"`
}

type ollamaOptions struct {
	NumPredict    int      `json:"num_predict"`
	Temperature   float64  `json:"temperature"`
	RepeatPenalty float64  `json:"repeat_penalty"`
	TopP          float64  `json:"top_p"`
	Stop          []string `json:"stop,omitempty"`
}

func (c *OllamaClient) body(req CompletionRequest) ollamaRequest {
	return ollamaRequest{
		Model:    c.model(req),
		Messages: req.Messages,
		Stream:   req.Stream,
		Options: ollamaOptions{
			NumPredict:    req.MaxTokens,
			Temperature:   req.Temperature,
			RepeatPenalty: req.RepetitionPenalty,
			TopP:          req.TopP,
			Stop:          req.Stop,
		},
	}
}

// Complete sends a non-streaming request and returns the full message content.
func (c *OllamaClient) Complete(ctx context.Context, req CompletionRequest) (string, error) {
	req.Stream = false
	ctx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()
	resp, err := c.post(ctx, c.body(req))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return "", httpError("ollama", resp)
	}

	var out struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return "", fmt.Errorf("ollama: decode response: %w", err)
	}
	return out.Message.Content, nil
}

// Stream sends a streaming request. Ollama emits newline-delimited JSON objects;
// we forward each message.content and stop when done:true.
func (c *OllamaClient) Stream(ctx context.Context, req CompletionRequest) (<-chan string, error) {
	req.Stream = true
	resp, err := c.post(ctx, c.body(req))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		defer resp.Body.Close()
		return nil, httpError("ollama", resp)
	}

	ch := make(chan string)
	go func() {
		defer close(ch)
		defer resp.Body.Close()
		sc := bufio.NewScanner(resp.Body)
		sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
		for sc.Scan() {
			line := sc.Bytes()
			if len(line) == 0 {
				continue
			}
			var chunk struct {
				Message struct {
					Content string `json:"content"`
				} `json:"message"`
				Done bool `json:"done"`
			}
			if err := json.Unmarshal(line, &chunk); err != nil {
				continue
			}
			if chunk.Message.Content != "" {
				select {
				case ch <- chunk.Message.Content:
				case <-ctx.Done():
					return
				}
			}
			if chunk.Done {
				return
			}
		}
	}()
	return ch, nil
}

func (c *OllamaClient) post(ctx context.Context, body any) (*http.Response, error) {
	return postJSON(ctx, c.endpoint+"/api/chat", body)
}