core/server/providers/openai.go

package providers

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

// OpenAIProvider handles any OpenAI-compatible API.
type OpenAIProvider struct{}

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

// ── Chat Completion (non-streaming) ─────────

func (p *OpenAIProvider) ChatCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (*CompletionResponse, error) {
	req.Stream = false

	body, err := p.doRequest(ctx, cfg, req)
	if err != nil {
		return nil, err
	}
	defer body.Close()

	var resp openaiChatResponse
	if err := json.NewDecoder(body).Decode(&resp); err != nil {
		return nil, fmt.Errorf("decode response: %w", err)
	}

	if len(resp.Choices) == 0 {
		return nil, fmt.Errorf("no choices in response")
	}

	return &CompletionResponse{
		Content:      resp.Choices[0].Message.Content,
		Model:        resp.Model,
		FinishReason: resp.Choices[0].FinishReason,
		InputTokens:  resp.Usage.PromptTokens,
		OutputTokens: resp.Usage.CompletionTokens,
	}, nil
}

// ── Stream Completion ───────────────────────

func (p *OpenAIProvider) StreamCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (<-chan StreamEvent, error) {
	req.Stream = true

	body, err := p.doRequest(ctx, cfg, req)
	if err != nil {
		return nil, err
	}

	ch := make(chan StreamEvent, 64)

	go func() {
		defer close(ch)
		defer body.Close()

		scanner := bufio.NewScanner(body)
		for scanner.Scan() {
			line := scanner.Text()

			if !strings.HasPrefix(line, "data: ") {
				continue
			}

			data := strings.TrimPrefix(line, "data: ")
			if data == "[DONE]" {
				ch <- StreamEvent{Done: true}
				return
			}

			var chunk openaiStreamChunk
			if err := json.Unmarshal([]byte(data), &chunk); err != nil {
				continue
			}

			if len(chunk.Choices) == 0 {
				continue
			}

			ev := StreamEvent{
				Delta:        chunk.Choices[0].Delta.Content,
				Model:        chunk.Model,
				FinishReason: chunk.Choices[0].FinishReason,
			}
			if ev.FinishReason != "" {
				ev.Done = true
			}

			ch <- ev
		}

		if err := scanner.Err(); err != nil {
			ch <- StreamEvent{Error: err}
		}
	}()

	return ch, nil
}

// ── List Models ─────────────────────────────

func (p *OpenAIProvider) ListModels(ctx context.Context, cfg ProviderConfig) ([]Model, error) {
	endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/models"

	httpReq, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil)
	if err != nil {
		return nil, err
	}
	if cfg.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}
	for k, v := range cfg.CustomHeaders {
		httpReq.Header.Set(k, v)
	}

	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("list models: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("list models: HTTP %d: %s", resp.StatusCode, string(b))
	}

	var result openaiModelsResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("decode models: %w", err)
	}

	models := make([]Model, 0, len(result.Data))
	for _, m := range result.Data {
		// Try known table first, then heuristic
		caps, found := LookupKnownModel(m.ID)
		if !found {
			caps = InferCapabilities(m.ID)
		}
		// Use context_length from API if available and we don't have it
		if m.ContextLength > 0 && caps.MaxContext == 0 {
			caps.MaxContext = m.ContextLength
		}

		models = append(models, Model{
			ID:           m.ID,
			OwnedBy:      m.OwnedBy,
			Capabilities: caps,
		})
	}
	return models, nil
}

// ── HTTP Layer ──────────────────────────────

func (p *OpenAIProvider) doRequest(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (io.ReadCloser, error) {
	endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/chat/completions"

	// Build OpenAI-format request
	oaiReq := openaiChatRequest{
		Model:    req.Model,
		Stream:   req.Stream,
		Messages: make([]openaiMessage, 0, len(req.Messages)),
	}
	if req.MaxTokens > 0 {
		oaiReq.MaxTokens = req.MaxTokens
	}
	if req.Temperature != nil {
		oaiReq.Temperature = req.Temperature
	}
	if req.TopP != nil {
		oaiReq.TopP = req.TopP
	}

	for _, m := range req.Messages {
		oaiReq.Messages = append(oaiReq.Messages, openaiMessage{
			Role:    m.Role,
			Content: m.Content,
		})
	}

	body, err := json.Marshal(oaiReq)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if cfg.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}
	// Inject provider-level custom headers (OpenRouter, Venice, etc.)
	for k, v := range cfg.CustomHeaders {
		httpReq.Header.Set(k, v)
	}

	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("provider request: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("provider error: HTTP %d: %s", resp.StatusCode, string(b))
	}

	return resp.Body, nil
}

// ── OpenAI Wire Types ───────────────────────

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

type openaiChatRequest struct {
	Model       string          `json:"model"`
	Messages    []openaiMessage `json:"messages"`
	MaxTokens   int             `json:"max_tokens,omitempty"`
	Temperature *float64        `json:"temperature,omitempty"`
	TopP        *float64        `json:"top_p,omitempty"`
	Stream      bool            `json:"stream"`
}

type openaiChatResponse struct {
	Model   string `json:"model"`
	Choices []struct {
		Message      openaiMessage `json:"message"`
		FinishReason string        `json:"finish_reason"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
	} `json:"usage"`
}

type openaiStreamChunk struct {
	Model   string `json:"model"`
	Choices []struct {
		Delta struct {
			Content string `json:"content"`
		} `json:"delta"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
}

type openaiModelsResponse struct {
	Data []struct {
		ID            string `json:"id"`
		OwnedBy       string `json:"owned_by"`
		ContextLength int    `json:"context_length,omitempty"` // Ollama, OpenRouter
	} `json:"data"`
}