GopherGate/internal/providers/openai.go

package providers

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/go-resty/resty/v2"
	"gophergate/internal/config"
	"gophergate/internal/models"
)

type OpenAIProvider struct {
	client *resty.Client
	config config.OpenAIConfig
	apiKey string
}

func NewOpenAIProvider(cfg config.OpenAIConfig, apiKey string) *OpenAIProvider {
	return &OpenAIProvider{
		client: resty.New().SetTimeout(10 * time.Minute),
		config: cfg,
		apiKey: apiKey,
	}
}

func (p *OpenAIProvider) Name() string {
	return "openai"
}

func (p *OpenAIProvider) ChatCompletion(ctx context.Context, req *models.UnifiedRequest) (*models.ChatCompletionResponse, error) {
	messagesJSON, err := MessagesToOpenAIJSON(req.Messages)
	if err != nil {
		return nil, fmt.Errorf("failed to convert messages: %w", err)
	}

	body := BuildOpenAIBody(req, messagesJSON, false)

	// Transition: Newer models require max_completion_tokens
	if strings.HasPrefix(req.Model, "o1-") || strings.HasPrefix(req.Model, "o3-") || strings.Contains(req.Model, "gpt-5") {
		if maxTokens, ok := body["max_tokens"]; ok {
			delete(body, "max_tokens")
			body["max_completion_tokens"] = maxTokens
		}
	}

	resp, err := p.client.R().
		SetContext(ctx).
		SetHeader("Authorization", "Bearer "+p.apiKey).
		SetBody(body).
		Post(fmt.Sprintf("%s/chat/completions", p.config.BaseURL))

	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}

	if !resp.IsSuccess() {
		return nil, fmt.Errorf("OpenAI API error (%d): %s", resp.StatusCode(), resp.String())
	}

	var respJSON map[string]interface{}
	if err := json.Unmarshal(resp.Body(), &respJSON); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	return ParseOpenAIResponse(respJSON, req.Model)
}

func (p *OpenAIProvider) ImageGeneration(ctx context.Context, req *models.ImageGenerationRequest) (*models.ImageGenerationResponse, error) {
	body := map[string]interface{}{
		"prompt": req.Prompt,
		"model":  req.Model,
	}

	if req.N != nil {
		body["n"] = *req.N
	}
	if req.Quality != nil {
		body["quality"] = *req.Quality
	}
	if req.ResponseFormat != nil {
		body["response_format"] = *req.ResponseFormat
	}
	if req.Size != nil {
		body["size"] = *req.Size
	}
	if req.Style != nil {
		body["style"] = *req.Style
	}
	if req.User != nil {
		body["user"] = *req.User
	}

	resp, err := p.client.R().
		SetContext(ctx).
		SetHeader("Authorization", "Bearer "+p.apiKey).
		SetBody(body).
		Post(fmt.Sprintf("%s/images/generations", p.config.BaseURL))

	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}

	if !resp.IsSuccess() {
		return nil, fmt.Errorf("OpenAI image API error (%d): %s", resp.StatusCode(), resp.String())
	}

	var result models.ImageGenerationResponse
	if err := json.Unmarshal(resp.Body(), &result); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	return &result, nil
}

func (p *OpenAIProvider) ChatCompletionStream(ctx context.Context, req *models.UnifiedRequest) (<-chan *models.ChatCompletionStreamResponse, error) {
	messagesJSON, err := MessagesToOpenAIJSON(req.Messages)
	if err != nil {
		return nil, fmt.Errorf("failed to convert messages: %w", err)
	}

	body := BuildOpenAIBody(req, messagesJSON, true)

	// Transition: Newer models require max_completion_tokens
	if strings.HasPrefix(req.Model, "o1-") || strings.HasPrefix(req.Model, "o3-") || strings.Contains(req.Model, "gpt-5") {
		if maxTokens, ok := body["max_tokens"]; ok {
			delete(body, "max_tokens")
			body["max_completion_tokens"] = maxTokens
		}
	}

	resp, err := p.client.R().
		SetContext(ctx).
		SetHeader("Authorization", "Bearer "+p.apiKey).
		SetBody(body).
		SetDoNotParseResponse(true).
		Post(fmt.Sprintf("%s/chat/completions", p.config.BaseURL))

	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}

	if !resp.IsSuccess() {
		return nil, fmt.Errorf("OpenAI API error (%d): %s", resp.StatusCode(), resp.String())
	}

	ch := make(chan *models.ChatCompletionStreamResponse)

	go func() {
		defer close(ch)
		err := StreamOpenAI(resp.RawBody(), ch)
		if err != nil {
			// In a real app, you might want to send an error chunk or log it
			fmt.Printf("Stream error: %v\n", err)
		}
	}()

	return ch, nil
}