naco-api/internal/ollama/client.go

package ollama

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/nacorid/naco-api/internal/clustering"
	"github.com/ollama/ollama/api"
)

const (
	// embedConcurrency caps how many embedding requests run in parallel. Ollama
	// serializes model requests server-side (OLLAMA_NUM_PARALLEL defaults to 1),
	// so unbounded fan-out would only add latency and memory for no throughput.
	embedConcurrency = 4

	// embedRetries and embedRetryBackoff cover transient server-side embedding
	// failures (e.g. cold model load, transient 5xx/429).
	embedRetries      = 2
	embedRetryBackoff = 150 * time.Millisecond

	// generateRetries and generateRetryBackoff cover transient failures that
	// produce no output at all; a truncated stream is never retried because the
	// partial response cannot be safely replayed.
	generateRetries      = 1
	generateRetryBackoff = 500 * time.Millisecond
)

type OllamaClient struct {
	client         api.Client
	embeddingModel string
	styleModel     string
}

func New(host *url.URL, httpClient *http.Client, embeddingModel, styleModel string) *OllamaClient {
	client := api.NewClient(host, httpClient)
	return &OllamaClient{client: *client, embeddingModel: embeddingModel, styleModel: styleModel}
}

// EmbedTexts computes one embedding per post. Embeddings are returned in the
// same order as posts; callers rely on index correspondence (e.g.
// items[i].Embedding = embeddings[i]) so posts are never skipped. Requests run
// in a bounded worker pool and transient failures are retried with backoff.
func (o *OllamaClient) EmbedTexts(ctx context.Context, posts []clustering.Item) ([][]float64, error) {
	if len(posts) == 0 {
		return nil, nil
	}

	embeddings := make([][]float64, len(posts))

	// poolCtx is a child of ctx: cancelling it on the first failure must not
	// cancel the caller's request context.
	poolCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	var (
		mu       sync.Mutex
		firstErr error
		sema     = make(chan struct{}, embedConcurrency)
		wg       sync.WaitGroup
	)
	for i, post := range posts {
		wg.Add(1)
		go func(idx int, text string) {
			defer wg.Done()
			select {
			case sema <- struct{}{}:
			case <-poolCtx.Done():
				return
			}
			defer func() { <-sema }()

			emb, err := o.embedWithRetry(poolCtx, strings.TrimSpace(text))
			if err != nil {
				mu.Lock()
				if firstErr == nil {
					firstErr = fmt.Errorf("embedding post %d (%q): %w", idx, text, err)
				}
				mu.Unlock()
				cancel()
				return
			}
			embeddings[idx] = emb
		}(i, post.Text)
	}
	wg.Wait()

	if firstErr := atomicErr(&mu, &firstErr); firstErr != nil {
		return embeddings, firstErr
	}
	// Every worker may have exited via the poolCtx.Done() branch without
	// recording an error, e.g. when the caller's context is already canceled.
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	return embeddings, nil
}

func atomicErr(mu *sync.Mutex, err *error) error {
	mu.Lock()
	defer mu.Unlock()
	return *err
}

func (o *OllamaClient) embedWithRetry(ctx context.Context, text string) ([]float64, error) {
	backoff := embedRetryBackoff
	var lastErr error
	for attempt := 0; attempt <= embedRetries; attempt++ {
		if attempt > 0 {
			if err := sleepCtx(ctx, backoff); err != nil {
				return nil, err
			}
			backoff *= 2
		}
		resp, err := o.client.Embeddings(ctx, &api.EmbeddingRequest{
			Model:  o.embeddingModel,
			Prompt: text,
		})
		if err == nil {
			return resp.Embedding, nil
		}
		lastErr = err
		if !isTransientErr(err) {
			return nil, err
		}
	}
	return nil, lastErr
}

func (o *OllamaClient) SummarizeStyle(ctx context.Context, posts []clustering.Item) (string, error) {
	if len(posts) == 0 {
		return "", nil
	}

	promptParts := make([]string, 0, len(posts))
	for _, post := range posts {
		if text := strings.TrimSpace(post.Text); text != "" {
			promptParts = append(promptParts, text)
		}
	}
	if len(promptParts) == 0 {
		return "", fmt.Errorf("no post text to summarize")
	}

	prompt := "Summarize the user's writing style. Analyze:\n" +
		"- tone and emotional register\n" +
		"- recurring themes and interests\n" +
		"- vocabulary and sentence structure\n" +
		"- communication habits\n" +
		"- personality traits expressed through writing\n\n" +
		"Base your analysis only on the supplied posts. Do not infer sensitive personal attributes.\n\n" +
		"Posts:\n\n" +
		strings.Join(promptParts, "\n---\n") +
		"\n\nStyle summary:"

	var fullResponse strings.Builder
	backoff := generateRetryBackoff
	var lastErr error
	for attempt := 0; attempt <= generateRetries; attempt++ {
		if attempt > 0 {
			if err := sleepCtx(ctx, backoff); err != nil {
				return "", err
			}
		}

		// Reset between attempts so a failed run never leaks a partial summary.
		fullResponse.Reset()
		done := false
		doneReason := ""
		err := o.client.Generate(ctx, &api.GenerateRequest{
			Model:  o.styleModel,
			Prompt: prompt,
		}, func(resp api.GenerateResponse) error {
			fullResponse.WriteString(resp.Response)
			if resp.Done {
				done = true
				doneReason = resp.DoneReason
			}
			return nil
		})
		if err != nil {
			lastErr = err
			// Only retry transport/5xx failures that produced no output; a
			// partial stream must not be replayed.
			if !isTransientErr(err) || fullResponse.Len() > 0 {
				return "", fmt.Errorf("ollama generate with model %q failed: %w", o.styleModel, err)
			}
			continue
		}

		// The underlying stream reader can finish without an error when the
		// request context is canceled or the connection drops mid-stream (the
		// library never inspects bufio.Scanner.Err), so a bare nil from
		// Generate does not mean the response is complete.
		if ctxErr := ctx.Err(); ctxErr != nil {
			return "", ctxErr
		}
		if !done {
			return "", fmt.Errorf("ollama stream for model %q ended before completion", o.styleModel)
		}
		if doneReason == "length" {
			return "", fmt.Errorf("ollama response for model %q was truncated (context limit reached)", o.styleModel)
		}
		return fullResponse.String(), nil
	}
	return "", fmt.Errorf("ollama generate with model %q failed after retries: %w", o.styleModel, lastErr)
}

func sleepCtx(ctx context.Context, d time.Duration) error {
	select {
	case <-time.After(d):
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// isTransientErr reports whether err is a retryable server-side failure. Non-2xx
// HTTP errors that carry a StatusCode (429, 5xx) qualify; context and transport
// errors below that never do, so they surface immediately.
func isTransientErr(err error) bool {
	var se api.StatusError
	if errors.As(err, &se) {
		return se.StatusCode == http.StatusTooManyRequests || se.StatusCode >= http.StatusInternalServerError
	}
	return false
}