package llm import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "strconv" "strings" "time" applogger "github.com/gochat/gochat/pkg/logger" ) // OpenAIProvider implements the Provider interface using OpenAI-compatible APIs. // Supports custom base_url for domestic providers like Volcengine/Doubao. type OpenAIProvider struct { apiKey string baseURL string model string embedModel string httpClient *http.Client maxRetries int } // OpenAIProviderConfig holds configuration for creating an OpenAIProvider. type OpenAIProviderConfig struct { APIKey string BaseURL string // defaults to "https://api.openai.com/v1" Model string // defaults to "gpt-4" EmbedModel string // defaults to "text-embedding-3-small" MaxRetries int // defaults to 3 Timeout int // HTTP timeout in seconds, defaults to 60 } // NewOpenAIProvider creates a new OpenAIProvider with the given configuration. func NewOpenAIProvider(cfg OpenAIProviderConfig) *OpenAIProvider { if cfg.BaseURL == "" { cfg.BaseURL = "https://api.openai.com/v1" } // Ensure baseURL ends without trailing slash cfg.BaseURL = strings.TrimRight(cfg.BaseURL, "/") if cfg.Model == "" { cfg.Model = "gpt-4" } if cfg.EmbedModel == "" { cfg.EmbedModel = "text-embedding-3-small" } if cfg.MaxRetries == 0 { cfg.MaxRetries = 3 } if cfg.Timeout == 0 { cfg.Timeout = 60 } return &OpenAIProvider{ apiKey: cfg.APIKey, baseURL: cfg.BaseURL, model: cfg.Model, embedModel: cfg.EmbedModel, httpClient: &http.Client{ Timeout: time.Duration(cfg.Timeout) * time.Second, }, maxRetries: cfg.MaxRetries, } } // ChatCompletion sends a chat completion request to the OpenAI-compatible API. func (p *OpenAIProvider) ChatCompletion(ctx context.Context, req ChatRequest) (*ChatResponse, error) { // Set default model if not specified if req.Model == "" { req.Model = p.model } body, err := json.Marshal(req) if err != nil { applogger.L().Errorf("ChatCompletion: failed to marshal request: %v", err) return nil, fmt.Errorf("marshal chat request: %w", err) } respBody, err := p.doRequestWithRetry(ctx, "/chat/completions", body) if err != nil { return nil, fmt.Errorf("chat completion request: %w", err) } var resp ChatResponse if err := json.Unmarshal(respBody, &resp); err != nil { applogger.L().Errorf("ChatCompletion: failed to unmarshal response: %v", err) return nil, fmt.Errorf("unmarshal chat response: %w", err) } return &resp, nil } // ChatCompletionStream sends a streaming chat completion request and returns chunks via callback. func (p *OpenAIProvider) ChatCompletionStream(ctx context.Context, req ChatRequest, onChunk func(StreamChunk) error) error { if req.Model == "" { req.Model = p.model } req.Stream = true body, err := json.Marshal(req) if err != nil { applogger.L().Errorf("ChatCompletionStream: failed to marshal request: %v", err) return fmt.Errorf("marshal chat request: %w", err) } httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/chat/completions", bytes.NewReader(body)) if err != nil { return fmt.Errorf("create stream request: %w", err) } p.setHeaders(httpReq) httpResp, err := p.httpClient.Do(httpReq) if err != nil { return fmt.Errorf("stream request: %w", err) } defer httpResp.Body.Close() if httpResp.StatusCode != http.StatusOK { respBody, _ := io.ReadAll(httpResp.Body) applogger.L().Errorf("ChatCompletionStream: unexpected status %d: %s", httpResp.StatusCode, string(respBody)) return fmt.Errorf("stream request status %d: %s", httpResp.StatusCode, string(respBody)) } return p.parseSSEStream(httpResp.Body, onChunk) } // CreateEmbedding sends an embedding request to the OpenAI-compatible API. func (p *OpenAIProvider) CreateEmbedding(ctx context.Context, req EmbeddingRequest) (*EmbeddingResponse, error) { if req.Model == "" { req.Model = p.embedModel } body, err := json.Marshal(req) if err != nil { applogger.L().Errorf("CreateEmbedding: failed to marshal request: %v", err) return nil, fmt.Errorf("marshal embedding request: %w", err) } respBody, err := p.doRequestWithRetry(ctx, "/embeddings", body) if err != nil { return nil, fmt.Errorf("embedding request: %w", err) } var resp EmbeddingResponse if err := json.Unmarshal(respBody, &resp); err != nil { applogger.L().Errorf("CreateEmbedding: failed to unmarshal response: %v", err) return nil, fmt.Errorf("unmarshal embedding response: %w", err) } return &resp, nil } // doRequestWithRetry performs an HTTP request with retry logic. func (p *OpenAIProvider) doRequestWithRetry(ctx context.Context, path string, body []byte) ([]byte, error) { var lastErr error for attempt := 0; attempt <= p.maxRetries; attempt++ { if attempt > 0 { // Exponential backoff: 1s, 2s, 4s backoff := time.Duration(1<= 400 { apiErr := parseAPIError(resp.StatusCode, respBody) applogger.L().Errorf("API error (status %d): %v", resp.StatusCode, apiErr) return nil, apiErr } return respBody, nil } // setHeaders sets common headers for API requests. func (p *OpenAIProvider) setHeaders(req *http.Request) { req.Header.Set("Content-Type", "application/json") req.Header.Set("Authorization", "Bearer "+p.apiKey) } // parseSSEStream parses Server-Sent Events from a streaming response body. func (p *OpenAIProvider) parseSSEStream(body io.Reader, onChunk func(StreamChunk) error) error { reader := newSSEReader(body) for { event, err := reader.Next() if err != nil { return fmt.Errorf("read SSE event: %w", err) } if event == nil { // Stream ended return nil } // Skip non-data events if event.Type != "message" || event.Data == "" { continue } // OpenAI sends "[DONE]" to signal stream end if event.Data == "[DONE]" { return nil } var chunk StreamChunk if err := json.Unmarshal([]byte(event.Data), &chunk); err != nil { applogger.L().Errorf("parseSSEStream: failed to unmarshal chunk: %v (data: %s)", err, event.Data) continue } if err := onChunk(chunk); err != nil { return fmt.Errorf("chunk callback: %w", err) } } } // --- Error types --- // APIError represents an error returned by the OpenAI-compatible API. type APIError struct { StatusCode int Message string Type string Code string } func (e *APIError) Error() string { return fmt.Sprintf("API error (status %d): %s", e.StatusCode, e.Message) } // parseAPIError creates an APIError from HTTP status and response body. func parseAPIError(statusCode int, body []byte) *APIError { apiErr := &APIError{ StatusCode: statusCode, Message: string(body), } // Try to parse OpenAI error structure var errResp struct { Error struct { Message string `json:"message"` Type string `json:"type"` Code string `json:"code"` } `json:"error"` } if err := json.Unmarshal(body, &errResp); err == nil && errResp.Error.Message != "" { apiErr.Message = errResp.Error.Message apiErr.Type = errResp.Error.Type apiErr.Code = errResp.Error.Code } return apiErr } // isNonRetriableError returns true for errors that should not be retried. func isNonRetriableError(err error) bool { if apiErr, ok := err.(*APIError); ok { // Retry on rate limit (429) and server errors (5xx) // Don't retry on client errors (400, 401, 403, 404, etc.) return apiErr.StatusCode >= 400 && apiErr.StatusCode < 500 && apiErr.StatusCode != 429 } return false } // --- SSE Reader --- // sseEvent represents a parsed SSE event. type sseEvent struct { Type string // event type (default "message" if not specified) Data string // data payload ID string // event ID } // sseReader reads Server-Sent Events from a stream. type sseReader struct { scanner *sseLineScanner } func newSSEReader(body io.Reader) *sseReader { return &sseReader{ scanner: newSSELineScanner(body), } } // Next reads the next SSE event from the stream. // Returns nil when the stream is complete. func (r *sseReader) Next() (*sseEvent, error) { var event *sseEvent for { line, err := r.scanner.Next() if err != nil { return nil, err } if line == nil { // Stream ended return event, nil } text := *line if text == "" { // Empty line = event boundary, dispatch current event if event != nil { return event, nil } continue } if strings.HasPrefix(text, ":") { // Comment, skip continue } field, value := parseSSEField(text) switch field { case "event": if event == nil { event = &sseEvent{} } event.Type = value case "data": if event == nil { event = &sseEvent{Type: "message"} } if event.Data != "" { event.Data += "\n" } event.Data += value case "id": if event == nil { event = &sseEvent{} } event.ID = value } } } func parseSSEField(line string) (field, value string) { idx := strings.Index(line, ":") if idx == -1 { return line, "" } field = line[:idx] value = strings.TrimLeft(line[idx+1:], " ") return field, value } // sseLineScanner reads lines from an SSE stream efficiently. type sseLineScanner struct { reader io.Reader buffer []byte hasData bool } func newSSELineScanner(reader io.Reader) *sseLineScanner { return &sseLineScanner{ reader: reader, buffer: make([]byte, 0, 4096), } } // Next returns the next line from the stream. // Returns nil when the stream is complete. func (s *sseLineScanner) Next() (*string, error) { for { // Check if we have a complete line in the buffer idx := bytes.IndexByte(s.buffer, '\n') if idx != -1 { line := string(s.buffer[:idx]) s.buffer = s.buffer[idx+1:] // Strip \r if present (CRLF) line = strings.TrimRight(line, "\r") return &line, nil } // Read more data tmp := make([]byte, 4096) n, err := s.reader.Read(tmp) if n > 0 { s.buffer = append(s.buffer, tmp[:n]...) s.hasData = true } if err != nil { if err == io.EOF { if len(s.buffer) > 0 { line := string(s.buffer) s.buffer = s.buffer[:0] line = strings.TrimRight(line, "\r") return &line, nil } return nil, nil } return nil, err } } } // --- Utility functions --- // ParseFloatEmbedding converts a slice of any (JSON numbers) to []float64. // Useful when embedding responses contain mixed numeric types. func ParseFloatEmbedding(raw []interface{}) []float64 { result := make([]float64, len(raw)) for i, v := range raw { switch n := v.(type) { case float64: result[i] = n case float32: result[i] = float64(n) case int: result[i] = float64(n) case int64: result[i] = float64(n) case string: f, err := strconv.ParseFloat(n, 64) if err == nil { result[i] = f } } } return result }