gh-commit/internal/ai/client.go

// Package ai talks to a local Ollama server. The rozoomcool SDK is used for
// client construction and model pulls; generation goes through a raw
// /api/generate POST because the SDK's Generate cannot pass think, options,
// or keep_alive — the levers this tool depends on for speed.
package ai

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

	ollama "github.com/rozoomcool/go-ollama-sdk"

	"github.com/prdlk/gh-commit/internal/diff"
	"github.com/prdlk/gh-commit/internal/ui"
)

// Generation profiles: small budget for one-line commit messages, larger
// budget and context for scope JSON over big file trees.
const (
	commitNumPredict = 96
	commitNumCtx     = 8192
	scopeNumPredict  = 1024
	scopeNumCtx      = 16384
	keepAlive        = "10m"
	temperature      = 0.2
	topP             = 0.9
)

// Client is a speed-tuned Ollama client for one model.
type Client struct {
	host    string
	model   string
	timeout time.Duration
	http    *http.Client
	sdk     *ollama.OllamaClient
}

// New builds a client for host/model with a per-request timeout.
func New(host, model string, timeout time.Duration) *Client {
	host = strings.TrimRight(host, "/")
	return &Client{
		host:    host,
		model:   model,
		timeout: timeout,
		http:    &http.Client{Timeout: timeout},
		sdk:     ollama.NewClient(host),
	}
}

// Model returns the configured model tag.
func (c *Client) Model() string { return c.model }

type generateOptions struct {
	Temperature float64 `json:"temperature"`
	TopP        float64 `json:"top_p"`
	NumPredict  int     `json:"num_predict"`
	NumCtx      int     `json:"num_ctx"`
}

type generateRequest struct {
	Model     string          `json:"model"`
	Prompt    string          `json:"prompt"`
	Stream    bool            `json:"stream"`
	Think     bool            `json:"think"`
	KeepAlive string          `json:"keep_alive"`
	Options   generateOptions `json:"options"`
}

type generateResponse struct {
	Response string `json:"response"`
	Error    string `json:"error"`
}

// generate POSTs a raw /api/generate request with the speed profile applied.
func (c *Client) generate(prompt string, numPredict, numCtx int) (string, error) {
	body, err := json.Marshal(generateRequest{
		Model:     c.model,
		Prompt:    prompt,
		Stream:    false,
		Think:     false,
		KeepAlive: keepAlive,
		Options: generateOptions{
			Temperature: temperature,
			TopP:        topP,
			NumPredict:  numPredict,
			NumCtx:      numCtx,
		},
	})
	if err != nil {
		return "", err
	}

	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.host+"/api/generate", bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.http.Do(req)
	if err != nil {
		return "", fmt.Errorf("ollama request failed: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}
	var out generateResponse
	if err := json.Unmarshal(data, &out); err != nil {
		return "", fmt.Errorf("ollama returned unexpected payload: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		if out.Error != "" {
			return "", fmt.Errorf("ollama: %s", out.Error)
		}
		return "", fmt.Errorf("ollama returned HTTP %d", resp.StatusCode)
	}
	// think:false is sent, but strip any reasoning block as defense in depth.
	return stripThink(out.Response), nil
}

// CommitMessage filters diff, prompts the model, and returns a cleaned
// commit message ("" when the model produced nothing usable).
func (c *Client) CommitMessage(rawDiff, scope string) (string, error) {
	prompt := buildCommitPrompt(diff.Filter(rawDiff), scope)
	out, err := c.generate(prompt, commitNumPredict, commitNumCtx)
	if err != nil {
		return "", err
	}
	return CleanCommitMessage(out), nil
}

// ScopesRaw prompts the model for a scope mapping and returns the raw
// response text (callers parse it so they can show diagnostics on failure).
func (c *Client) ScopesRaw(filetree string, existing map[string][]string) (string, error) {
	return c.generate(buildScopePrompt(filetree, existing), scopeNumPredict, scopeNumCtx)
}

type tagsResponse struct {
	Models []struct {
		Name string `json:"name"`
	} `json:"models"`
}

// EnsureReady verifies the Ollama server is answering and the model is
// available locally, offering to pull it when missing.
func (c *Client) EnsureReady() error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.host+"/api/tags", nil)
	if err != nil {
		return err
	}
	resp, err := c.http.Do(req)
	if err != nil {
		return fmt.Errorf("the Ollama server is not reachable at %s — start it with 'ollama serve'", c.host)
	}
	defer func() { _ = resp.Body.Close() }()

	var tags tagsResponse
	if err := json.NewDecoder(resp.Body).Decode(&tags); err != nil {
		return fmt.Errorf("the Ollama server at %s gave an unexpected response — start it with 'ollama serve'", c.host)
	}
	for _, m := range tags.Models {
		if m.Name == c.model || (!strings.Contains(c.model, ":") && m.Name == c.model+":latest") {
			return nil
		}
	}

	ui.Warnf("Model %s is not available locally", c.model)
	if !ui.Confirm(fmt.Sprintf("Pull %s now?", c.model)) {
		return fmt.Errorf("model %s is not available — run 'ollama pull %s'", c.model, c.model)
	}
	last := ""
	if err := c.sdk.PullModel(c.model, func(status string) {
		if status != "" && status != last {
			ui.Dimf("  %s", status)
			last = status
		}
	}); err != nil {
		return fmt.Errorf("pulling %s: %w", c.model, err)
	}
	ui.Successf("✓ Pulled %s", c.model)
	return nil
}