descrybe/apps/api/internal/processing/openai.go

package processing

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/descrybe/descrybe-v2/apps/api/internal/config"
	"github.com/descrybe/descrybe-v2/apps/api/internal/security"
)

// OpenAIClient calls an OpenAI-compatible Chat Completions API with rate limiting and retries.
type OpenAIClient struct {
	APIKey      string
	BaseURL     string
	Model       string
	HTTPClient  *http.Client
	MinInterval time.Duration
	MaxRetries  int
	// ModeLabel is recorded on products/jobs for analytics
	// ("internal" | "popular:<name>" | "custom"). Defaults to internal.
	ModeLabel string

	mu       sync.Mutex
	lastCall time.Time
}

const maxOpenAIRetries = 8

func NewOpenAIClient(apiKey, baseURL, model string, rpm, maxRetries int) *OpenAIClient {
	if baseURL == "" {
		baseURL = "https://api.openai.com/v1"
	}
	if model == "" {
		model = "gpt-4o-mini"
	}
	if maxRetries <= 0 {
		maxRetries = 3
	} else if maxRetries > maxOpenAIRetries {
		maxRetries = maxOpenAIRetries
	}
	interval := time.Duration(0)
	if rpm > 0 {
		interval = time.Minute / time.Duration(rpm)
	}
	// Dial-time SSRF. Loopback when base URL is local; RFC1918 only in non-prod.
	policy := openAIDialPolicy(baseURL)
	return &OpenAIClient{
		APIKey:      apiKey,
		BaseURL:     strings.TrimRight(baseURL, "/"),
		Model:       model,
		HTTPClient:  security.SafeHTTPClientPolicy(60*time.Second, policy),
		MinInterval: interval,
		MaxRetries:  maxRetries,
		ModeLabel:   AIProviderInternal,
	}
}

func openAIDialPolicy(baseURL string) security.DialPolicy {
	if openAIBaseAllowsLoopback(baseURL) {
		return security.DialPolicy{AllowLoopback: true}
	}
	if openAIBaseAllowsPrivate(baseURL) {
		return security.DialPolicy{AllowLoopback: true, AllowPrivate: true}
	}
	return security.DialPolicy{}
}

func openAIBaseAllowsLoopback(baseURL string) bool {
	u, err := url.Parse(baseURL)
	if err != nil || u.Hostname() == "" {
		return false
	}
	host := strings.ToLower(u.Hostname())
	if host == "localhost" {
		return true
	}
	ip := net.ParseIP(host)
	return ip != nil && ip.IsLoopback()
}

// openAIBaseAllowsPrivate permits RFC1918/ULA literal OPENAI_BASE_URL hosts
// when APP_ENV is not production/prod (local LAN OpenAI-compatible proxies).
func openAIBaseAllowsPrivate(baseURL string) bool {
	if config.IsProductionEnv() {
		return false
	}
	u, err := url.Parse(baseURL)
	if err != nil || u.Hostname() == "" {
		return false
	}
	ip := net.ParseIP(strings.ToLower(u.Hostname()))
	if ip == nil || ip.IsLinkLocalUnicast() {
		return false
	}
	return ip.IsPrivate()
}

func (c *OpenAIClient) Enabled() bool {
	return c != nil && strings.TrimSpace(c.APIKey) != ""
}

// ProviderModeLabel implements ProviderLabeler for analytics writes.
func (c *OpenAIClient) ProviderModeLabel() string {
	if c == nil {
		return AIProviderUnknown
	}
	if label := strings.TrimSpace(c.ModeLabel); label != "" {
		return label
	}
	return AIProviderInternal
}

type chatRequest struct {
	Model       string        `json:"model"`
	Messages    []chatMessage `json:"messages"`
	Temperature float64       `json:"temperature"`
	MaxTokens   int           `json:"max_tokens,omitempty"`
}

type chatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type chatResponse struct {
	Model   string `json:"model"`
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
	Error *struct {
		Message string `json:"message"`
		Type    string `json:"type"`
	} `json:"error"`
}

func (c *OpenAIClient) Complete(ctx context.Context, system, user string) (Completion, error) {
	return c.CompleteWithOptions(ctx, system, user, CompleteOptions{})
}

func (c *OpenAIClient) CompleteWithOptions(ctx context.Context, system, user string, opts CompleteOptions) (Completion, error) {
	if !c.Enabled() {
		return Completion{}, errors.New("openai api key not configured")
	}
	system = SanitizeText(system)
	user = SanitizeText(user)
	temp := opts.Temperature
	if temp <= 0 {
		temp = DefaultStructuredTemp
	}
	if temp > 0.3 {
		temp = 0.3
	}

	var lastErr error
	for attempt := 0; attempt <= c.MaxRetries; attempt++ {
		if attempt > 0 {
			backoff := time.Duration(math.Pow(2, float64(attempt-1))) * 200 * time.Millisecond
			jitter := time.Duration(rand.Intn(100)) * time.Millisecond
			select {
			case <-ctx.Done():
				return Completion{}, ctx.Err()
			case <-time.After(backoff + jitter):
			}
		}
		if err := c.waitRate(ctx); err != nil {
			return Completion{}, err
		}
		comp, retryable, err := c.doComplete(ctx, system, user, temp, opts.MaxTokens)
		if err == nil {
			return comp, nil
		}
		lastErr = err
		if !retryable {
			return Completion{}, err
		}
	}
	return Completion{}, fmt.Errorf("openai retries exhausted: %w", lastErr)
}

func (c *OpenAIClient) waitRate(ctx context.Context) error {
	c.mu.Lock()
	if c.MinInterval <= 0 {
		c.lastCall = time.Now()
		c.mu.Unlock()
		return nil
	}
	now := time.Now()
	wait := c.MinInterval - now.Sub(c.lastCall)
	if wait < 0 {
		wait = 0
	}
	// Reserve the next slot under the lock so concurrent callers cannot both
	// observe the same lastCall and bypass MinInterval.
	c.lastCall = now.Add(wait)
	c.mu.Unlock()
	if wait > 0 {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(wait):
		}
	}
	return nil
}

type embeddingRequest struct {
	Model string   `json:"model"`
	Input []string `json:"input"`
}

type embeddingResponse struct {
	Data []struct {
		Embedding []float32 `json:"embedding"`
		Index     int       `json:"index"`
	} `json:"data"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

// Embed implements Embedder via OpenAI-compatible POST /embeddings.
func (c *OpenAIClient) Embed(ctx context.Context, texts []string) ([][]float32, error) {
	if !c.Enabled() {
		return nil, errors.New("openai api key not configured")
	}
	if len(texts) == 0 {
		return nil, errors.New("empty embedding input")
	}
	clean := make([]string, 0, len(texts))
	for _, t := range texts {
		t = SanitizeText(t)
		if t == "" {
			return nil, errors.New("empty embedding input")
		}
		clean = append(clean, t)
	}

	var lastErr error
	for attempt := 0; attempt <= c.MaxRetries; attempt++ {
		if attempt > 0 {
			backoff := time.Duration(math.Pow(2, float64(attempt-1))) * 200 * time.Millisecond
			jitter := time.Duration(rand.Intn(100)) * time.Millisecond
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(backoff + jitter):
			}
		}
		if err := c.waitRate(ctx); err != nil {
			return nil, err
		}
		vecs, retryable, err := c.doEmbed(ctx, clean)
		if err == nil {
			return vecs, nil
		}
		lastErr = err
		if !retryable {
			return nil, err
		}
	}
	return nil, fmt.Errorf("openai embedding retries exhausted: %w", lastErr)
}

func (c *OpenAIClient) doEmbed(ctx context.Context, texts []string) ([][]float32, bool, error) {
	body, err := json.Marshal(embeddingRequest{Model: c.Model, Input: texts})
	if err != nil {
		return nil, false, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/embeddings", bytes.NewReader(body))
	if err != nil {
		return nil, false, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.APIKey)

	res, err := c.HTTPClient.Do(req)
	if err != nil {
		return nil, true, err
	}
	defer res.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(res.Body, 4<<20))
	if err != nil {
		return nil, true, err
	}
	var parsed embeddingResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return nil, false, fmt.Errorf("openai embeddings decode: %w", err)
	}
	if res.StatusCode == http.StatusTooManyRequests || res.StatusCode >= 500 {
		msg := "rate limited or server error"
		if parsed.Error != nil && parsed.Error.Message != "" {
			msg = TruncateError(errors.New(parsed.Error.Message))
		}
		return nil, true, errors.New(msg)
	}
	if res.StatusCode >= 400 {
		msg := fmt.Sprintf("openai embeddings http %d", res.StatusCode)
		if parsed.Error != nil && parsed.Error.Message != "" {
			msg = TruncateError(errors.New(parsed.Error.Message))
		}
		return nil, false, errors.New(msg)
	}
	if len(parsed.Data) == 0 {
		return nil, false, errors.New("empty embedding response")
	}
	out := make([][]float32, len(texts))
	for _, row := range parsed.Data {
		if row.Index < 0 || row.Index >= len(out) {
			return nil, false, errors.New("embedding index out of range")
		}
		out[row.Index] = row.Embedding
	}
	for i, v := range out {
		if len(v) == 0 {
			return nil, false, fmt.Errorf("missing embedding at index %d", i)
		}
	}
	return out, false, nil
}

func (c *OpenAIClient) doComplete(ctx context.Context, system, user string, temperature float64, maxTokens int) (Completion, bool, error) {
	reqBody := chatRequest{
		Model: c.Model,
		Messages: []chatMessage{
			{Role: "system", Content: system},
			{Role: "user", Content: user},
		},
		Temperature: temperature,
		MaxTokens:   maxTokens,
	}
	body, err := json.Marshal(reqBody)
	if err != nil {
		return Completion{}, false, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return Completion{}, false, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.APIKey)

	res, err := c.HTTPClient.Do(req)
	if err != nil {
		return Completion{}, true, err
	}
	defer res.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(res.Body, 1<<20))
	if err != nil {
		return Completion{}, true, err
	}
	var parsed chatResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return Completion{}, false, fmt.Errorf("openai decode: %w", err)
	}
	if res.StatusCode == http.StatusTooManyRequests || res.StatusCode >= 500 {
		msg := "rate limited or server error"
		if parsed.Error != nil && parsed.Error.Message != "" {
			msg = TruncateError(errors.New(parsed.Error.Message))
		}
		return Completion{}, true, errors.New(msg)
	}
	if res.StatusCode >= 400 {
		msg := fmt.Sprintf("openai http %d", res.StatusCode)
		if parsed.Error != nil && parsed.Error.Message != "" {
			msg = TruncateError(errors.New(parsed.Error.Message))
		}
		return Completion{}, false, errors.New(msg)
	}
	text := ""
	if len(parsed.Choices) > 0 {
		text = SanitizeOutput(parsed.Choices[0].Message.Content)
	}
	if text == "" {
		return Completion{}, false, errors.New("empty model response")
	}
	return Completion{
		Text:         text,
		PromptTokens: parsed.Usage.PromptTokens,
		OutputTokens: parsed.Usage.CompletionTokens,
		TotalTokens:  parsed.Usage.TotalTokens,
		Model:        parsed.Model,
		Raw: map[string]any{
			"model":  parsed.Model,
			"usage":  parsed.Usage,
			"status": res.StatusCode,
		},
	}, false, nil
}

// HeuristicCompleter is used when OpenAI is not configured (local/dev fallback).
type HeuristicCompleter struct{}

// ProviderModeLabel labels heuristic output for analytics (not a paid provider).
func (h HeuristicCompleter) ProviderModeLabel() string {
	return AIProviderInternal
}

func (h HeuristicCompleter) Complete(_ context.Context, system, user string) (Completion, error) {
	systemL := strings.ToLower(system)
	user = SanitizeText(user)
	text := "General"
	switch {
	case strings.Contains(systemL, `"name"`) || strings.Contains(systemL, "titles and descriptions"):
		// Prefer explicit Name:/Desc: (ProductEnhanceUser) or Current name: labels.
		// Never use firstLine(user) alone — that line is often "Category: …".
		name := labeledPromptValue(user, "name:", "current name:")
		if name == "" || isPromptLabelTitle(name) {
			name = "Product"
		}
		desc := labeledPromptValue(user, "desc:", "description:", "current description:")
		if desc == "" {
			desc = "Product description"
		}
		b, _ := json.Marshal(map[string]string{"name": name, "description": desc})
		text = string(b)
	case strings.Contains(systemL, "attributes") && strings.Contains(systemL, "json"):
		text = `{"material":"unknown","brand":"unknown"}`
	case strings.Contains(systemL, "categor"):
		text = "General"
	default:
		// Never echo ProductEnhanceUser's first line ("Category: …") as title/output.
		text = labeledPromptValue(user, "name:", "current name:")
		if text == "" || isPromptLabelTitle(text) {
			text = "ok"
		}
	}
	return Completion{
		Text:        text,
		TotalTokens: 0,
		Model:       "heuristic",
		Raw:         map[string]any{"provider": "heuristic"},
	}, nil
}