1package provider
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"errors"
  7	"fmt"
  8	"io"
  9	"log/slog"
 10	"time"
 11
 12	"github.com/charmbracelet/crush/internal/config"
 13	"github.com/charmbracelet/crush/internal/fur/provider"
 14	"github.com/charmbracelet/crush/internal/llm/tools"
 15	"github.com/charmbracelet/crush/internal/message"
 16	"github.com/openai/openai-go"
 17	"github.com/openai/openai-go/option"
 18	"github.com/openai/openai-go/shared"
 19)
 20
 21type openaiClient struct {
 22	providerOptions providerClientOptions
 23	client          openai.Client
 24}
 25
 26type OpenAIClient ProviderClient
 27
 28func newOpenAIClient(opts providerClientOptions) OpenAIClient {
 29	return &openaiClient{
 30		providerOptions: opts,
 31		client:          createOpenAIClient(opts),
 32	}
 33}
 34
 35func createOpenAIClient(opts providerClientOptions) openai.Client {
 36	openaiClientOptions := []option.RequestOption{}
 37	if opts.apiKey != "" {
 38		openaiClientOptions = append(openaiClientOptions, option.WithAPIKey(opts.apiKey))
 39	}
 40	if opts.baseURL != "" {
 41		resolvedBaseURL, err := config.Get().Resolve(opts.baseURL)
 42		if err == nil {
 43			openaiClientOptions = append(openaiClientOptions, option.WithBaseURL(resolvedBaseURL))
 44		}
 45	}
 46
 47	if opts.extraHeaders != nil {
 48		for key, value := range opts.extraHeaders {
 49			openaiClientOptions = append(openaiClientOptions, option.WithHeader(key, value))
 50		}
 51	}
 52
 53	return openai.NewClient(openaiClientOptions...)
 54}
 55
 56func (o *openaiClient) convertMessages(messages []message.Message) (openaiMessages []openai.ChatCompletionMessageParamUnion) {
 57	// Add system message first
 58	openaiMessages = append(openaiMessages, openai.SystemMessage(o.providerOptions.systemMessage))
 59
 60	for _, msg := range messages {
 61		switch msg.Role {
 62		case message.User:
 63			var content []openai.ChatCompletionContentPartUnionParam
 64			textBlock := openai.ChatCompletionContentPartTextParam{Text: msg.Content().String()}
 65			content = append(content, openai.ChatCompletionContentPartUnionParam{OfText: &textBlock})
 66			for _, binaryContent := range msg.BinaryContent() {
 67				imageURL := openai.ChatCompletionContentPartImageImageURLParam{URL: binaryContent.String(provider.InferenceProviderOpenAI)}
 68				imageBlock := openai.ChatCompletionContentPartImageParam{ImageURL: imageURL}
 69
 70				content = append(content, openai.ChatCompletionContentPartUnionParam{OfImageURL: &imageBlock})
 71			}
 72
 73			openaiMessages = append(openaiMessages, openai.UserMessage(content))
 74
 75		case message.Assistant:
 76			assistantMsg := openai.ChatCompletionAssistantMessageParam{
 77				Role: "assistant",
 78			}
 79
 80			if msg.Content().String() != "" {
 81				assistantMsg.Content = openai.ChatCompletionAssistantMessageParamContentUnion{
 82					OfString: openai.String(msg.Content().String()),
 83				}
 84			}
 85
 86			if len(msg.ToolCalls()) > 0 {
 87				assistantMsg.ToolCalls = make([]openai.ChatCompletionMessageToolCallParam, len(msg.ToolCalls()))
 88				for i, call := range msg.ToolCalls() {
 89					assistantMsg.ToolCalls[i] = openai.ChatCompletionMessageToolCallParam{
 90						ID:   call.ID,
 91						Type: "function",
 92						Function: openai.ChatCompletionMessageToolCallFunctionParam{
 93							Name:      call.Name,
 94							Arguments: call.Input,
 95						},
 96					}
 97				}
 98			}
 99
100			openaiMessages = append(openaiMessages, openai.ChatCompletionMessageParamUnion{
101				OfAssistant: &assistantMsg,
102			})
103
104		case message.Tool:
105			for _, result := range msg.ToolResults() {
106				openaiMessages = append(openaiMessages,
107					openai.ToolMessage(result.Content, result.ToolCallID),
108				)
109			}
110		}
111	}
112
113	return
114}
115
116func (o *openaiClient) convertTools(tools []tools.BaseTool) []openai.ChatCompletionToolParam {
117	openaiTools := make([]openai.ChatCompletionToolParam, len(tools))
118
119	for i, tool := range tools {
120		info := tool.Info()
121		openaiTools[i] = openai.ChatCompletionToolParam{
122			Function: openai.FunctionDefinitionParam{
123				Name:        info.Name,
124				Description: openai.String(info.Description),
125				Parameters: openai.FunctionParameters{
126					"type":       "object",
127					"properties": info.Parameters,
128					"required":   info.Required,
129				},
130			},
131		}
132	}
133
134	return openaiTools
135}
136
137func (o *openaiClient) finishReason(reason string) message.FinishReason {
138	switch reason {
139	case "stop":
140		return message.FinishReasonEndTurn
141	case "length":
142		return message.FinishReasonMaxTokens
143	case "tool_calls":
144		return message.FinishReasonToolUse
145	default:
146		return message.FinishReasonUnknown
147	}
148}
149
150func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessageParamUnion, tools []openai.ChatCompletionToolParam) openai.ChatCompletionNewParams {
151	model := o.providerOptions.model(o.providerOptions.modelType)
152	cfg := config.Get()
153
154	modelConfig := cfg.Models[config.SelectedModelTypeLarge]
155	if o.providerOptions.modelType == config.SelectedModelTypeSmall {
156		modelConfig = cfg.Models[config.SelectedModelTypeSmall]
157	}
158
159	reasoningEffort := modelConfig.ReasoningEffort
160
161	params := openai.ChatCompletionNewParams{
162		Model:    openai.ChatModel(model.ID),
163		Messages: messages,
164		Tools:    tools,
165	}
166
167	maxTokens := model.DefaultMaxTokens
168	if modelConfig.MaxTokens > 0 {
169		maxTokens = modelConfig.MaxTokens
170	}
171
172	// Override max tokens if set in provider options
173	if o.providerOptions.maxTokens > 0 {
174		maxTokens = o.providerOptions.maxTokens
175	}
176	if model.CanReason {
177		params.MaxCompletionTokens = openai.Int(maxTokens)
178		switch reasoningEffort {
179		case "low":
180			params.ReasoningEffort = shared.ReasoningEffortLow
181		case "medium":
182			params.ReasoningEffort = shared.ReasoningEffortMedium
183		case "high":
184			params.ReasoningEffort = shared.ReasoningEffortHigh
185		default:
186			params.ReasoningEffort = shared.ReasoningEffort(reasoningEffort)
187		}
188	} else {
189		params.MaxTokens = openai.Int(maxTokens)
190	}
191
192	return params
193}
194
195func (o *openaiClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (response *ProviderResponse, err error) {
196	params := o.preparedParams(o.convertMessages(messages), o.convertTools(tools))
197	cfg := config.Get()
198	if cfg.Options.Debug {
199		jsonData, _ := json.Marshal(params)
200		slog.Debug("Prepared messages", "messages", string(jsonData))
201	}
202	attempts := 0
203	for {
204		attempts++
205		openaiResponse, err := o.client.Chat.Completions.New(
206			ctx,
207			params,
208		)
209		// If there is an error we are going to see if we can retry the call
210		if err != nil {
211			retry, after, retryErr := o.shouldRetry(attempts, err)
212			if retryErr != nil {
213				return nil, retryErr
214			}
215			if retry {
216				slog.Warn(fmt.Sprintf("Retrying due to rate limit... attempt %d of %d", attempts, maxRetries))
217				select {
218				case <-ctx.Done():
219					return nil, ctx.Err()
220				case <-time.After(time.Duration(after) * time.Millisecond):
221					continue
222				}
223			}
224			return nil, retryErr
225		}
226
227		content := ""
228		if openaiResponse.Choices[0].Message.Content != "" {
229			content = openaiResponse.Choices[0].Message.Content
230		}
231
232		toolCalls := o.toolCalls(*openaiResponse)
233		finishReason := o.finishReason(string(openaiResponse.Choices[0].FinishReason))
234
235		if len(toolCalls) > 0 {
236			finishReason = message.FinishReasonToolUse
237		}
238
239		return &ProviderResponse{
240			Content:      content,
241			ToolCalls:    toolCalls,
242			Usage:        o.usage(*openaiResponse),
243			FinishReason: finishReason,
244		}, nil
245	}
246}
247
248func (o *openaiClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {
249	params := o.preparedParams(o.convertMessages(messages), o.convertTools(tools))
250	params.StreamOptions = openai.ChatCompletionStreamOptionsParam{
251		IncludeUsage: openai.Bool(true),
252	}
253
254	cfg := config.Get()
255	if cfg.Options.Debug {
256		jsonData, _ := json.Marshal(params)
257		slog.Debug("Prepared messages", "messages", string(jsonData))
258	}
259
260	attempts := 0
261	eventChan := make(chan ProviderEvent)
262
263	go func() {
264		for {
265			attempts++
266			openaiStream := o.client.Chat.Completions.NewStreaming(
267				ctx,
268				params,
269			)
270
271			acc := openai.ChatCompletionAccumulator{}
272			currentContent := ""
273			toolCalls := make([]message.ToolCall, 0)
274
275			for openaiStream.Next() {
276				chunk := openaiStream.Current()
277				acc.AddChunk(chunk)
278
279				for _, choice := range chunk.Choices {
280					if choice.Delta.Content != "" {
281						eventChan <- ProviderEvent{
282							Type:    EventContentDelta,
283							Content: choice.Delta.Content,
284						}
285						currentContent += choice.Delta.Content
286					}
287				}
288			}
289
290			err := openaiStream.Err()
291			if err == nil || errors.Is(err, io.EOF) {
292				if cfg.Options.Debug {
293					jsonData, _ := json.Marshal(acc.ChatCompletion)
294					slog.Debug("Response", "messages", string(jsonData))
295				}
296				resultFinishReason := acc.ChatCompletion.Choices[0].FinishReason
297				if resultFinishReason == "" {
298					// If the finish reason is empty, we assume it was a successful completion
299					// INFO: this is happening for openrouter for some reason
300					resultFinishReason = "stop"
301				}
302				// Stream completed successfully
303				finishReason := o.finishReason(resultFinishReason)
304				if len(acc.Choices[0].Message.ToolCalls) > 0 {
305					toolCalls = append(toolCalls, o.toolCalls(acc.ChatCompletion)...)
306				}
307				if len(toolCalls) > 0 {
308					finishReason = message.FinishReasonToolUse
309				}
310
311				eventChan <- ProviderEvent{
312					Type: EventComplete,
313					Response: &ProviderResponse{
314						Content:      currentContent,
315						ToolCalls:    toolCalls,
316						Usage:        o.usage(acc.ChatCompletion),
317						FinishReason: finishReason,
318					},
319				}
320				close(eventChan)
321				return
322			}
323
324			// If there is an error we are going to see if we can retry the call
325			retry, after, retryErr := o.shouldRetry(attempts, err)
326			if retryErr != nil {
327				eventChan <- ProviderEvent{Type: EventError, Error: retryErr}
328				close(eventChan)
329				return
330			}
331			if retry {
332				slog.Warn(fmt.Sprintf("Retrying due to rate limit... attempt %d of %d", attempts, maxRetries))
333				select {
334				case <-ctx.Done():
335					// context cancelled
336					if ctx.Err() == nil {
337						eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
338					}
339					close(eventChan)
340					return
341				case <-time.After(time.Duration(after) * time.Millisecond):
342					continue
343				}
344			}
345			eventChan <- ProviderEvent{Type: EventError, Error: retryErr}
346			close(eventChan)
347			return
348		}
349	}()
350
351	return eventChan
352}
353
354func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error) {
355	var apiErr *openai.Error
356	if !errors.As(err, &apiErr) {
357		return false, 0, err
358	}
359
360	if attempts > maxRetries {
361		return false, 0, fmt.Errorf("maximum retry attempts reached for rate limit: %d retries", maxRetries)
362	}
363
364	// Check for token expiration (401 Unauthorized)
365	if apiErr.StatusCode == 401 {
366		o.providerOptions.apiKey, err = config.Get().Resolve(o.providerOptions.config.APIKey)
367		if err != nil {
368			return false, 0, fmt.Errorf("failed to resolve API key: %w", err)
369		}
370		o.client = createOpenAIClient(o.providerOptions)
371		return true, 0, nil
372	}
373
374	if apiErr.StatusCode != 429 && apiErr.StatusCode != 500 {
375		return false, 0, err
376	}
377
378	retryMs := 0
379	retryAfterValues := apiErr.Response.Header.Values("Retry-After")
380
381	backoffMs := 2000 * (1 << (attempts - 1))
382	jitterMs := int(float64(backoffMs) * 0.2)
383	retryMs = backoffMs + jitterMs
384	if len(retryAfterValues) > 0 {
385		if _, err := fmt.Sscanf(retryAfterValues[0], "%d", &retryMs); err == nil {
386			retryMs = retryMs * 1000
387		}
388	}
389	return true, int64(retryMs), nil
390}
391
392func (o *openaiClient) toolCalls(completion openai.ChatCompletion) []message.ToolCall {
393	var toolCalls []message.ToolCall
394
395	if len(completion.Choices) > 0 && len(completion.Choices[0].Message.ToolCalls) > 0 {
396		for _, call := range completion.Choices[0].Message.ToolCalls {
397			toolCall := message.ToolCall{
398				ID:       call.ID,
399				Name:     call.Function.Name,
400				Input:    call.Function.Arguments,
401				Type:     "function",
402				Finished: true,
403			}
404			toolCalls = append(toolCalls, toolCall)
405		}
406	}
407
408	return toolCalls
409}
410
411func (o *openaiClient) usage(completion openai.ChatCompletion) TokenUsage {
412	cachedTokens := completion.Usage.PromptTokensDetails.CachedTokens
413	inputTokens := completion.Usage.PromptTokens - cachedTokens
414
415	return TokenUsage{
416		InputTokens:         inputTokens,
417		OutputTokens:        completion.Usage.CompletionTokens,
418		CacheCreationTokens: 0, // OpenAI doesn't provide this directly
419		CacheReadTokens:     cachedTokens,
420	}
421}
422
423func (a *openaiClient) Model() provider.Model {
424	return a.providerOptions.model(a.providerOptions.modelType)
425}