openai.go

  1package provider
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"errors"
  7	"fmt"
  8	"io"
  9	"log/slog"
 10	"slices"
 11	"strings"
 12	"time"
 13
 14	"github.com/charmbracelet/catwalk/pkg/catwalk"
 15	"github.com/charmbracelet/crush/internal/config"
 16	"github.com/charmbracelet/crush/internal/llm/tools"
 17	"github.com/charmbracelet/crush/internal/log"
 18	"github.com/charmbracelet/crush/internal/message"
 19	"github.com/google/uuid"
 20	"github.com/openai/openai-go"
 21	"github.com/openai/openai-go/option"
 22	"github.com/openai/openai-go/packages/param"
 23	"github.com/openai/openai-go/shared"
 24)
 25
 26type openaiClient struct {
 27	providerOptions providerClientOptions
 28	client          openai.Client
 29}
 30
 31type OpenAIClient ProviderClient
 32
 33func newOpenAIClient(opts providerClientOptions) OpenAIClient {
 34	return &openaiClient{
 35		providerOptions: opts,
 36		client:          createOpenAIClient(opts),
 37	}
 38}
 39
 40func createOpenAIClient(opts providerClientOptions) openai.Client {
 41	openaiClientOptions := []option.RequestOption{}
 42	if opts.apiKey != "" {
 43		openaiClientOptions = append(openaiClientOptions, option.WithAPIKey(opts.apiKey))
 44	}
 45	if opts.baseURL != "" {
 46		resolvedBaseURL, err := config.Get().Resolve(opts.baseURL)
 47		if err == nil {
 48			openaiClientOptions = append(openaiClientOptions, option.WithBaseURL(resolvedBaseURL))
 49		}
 50	}
 51
 52	if config.Get().Options.Debug {
 53		httpClient := log.NewHTTPClient()
 54		openaiClientOptions = append(openaiClientOptions, option.WithHTTPClient(httpClient))
 55	}
 56
 57	for key, value := range opts.extraHeaders {
 58		openaiClientOptions = append(openaiClientOptions, option.WithHeader(key, value))
 59	}
 60
 61	for extraKey, extraValue := range opts.extraBody {
 62		openaiClientOptions = append(openaiClientOptions, option.WithJSONSet(extraKey, extraValue))
 63	}
 64
 65	return openai.NewClient(openaiClientOptions...)
 66}
 67
 68func (o *openaiClient) convertMessages(messages []message.Message) (openaiMessages []openai.ChatCompletionMessageParamUnion) {
 69	isAnthropicModel := o.providerOptions.config.ID == string(catwalk.InferenceProviderOpenRouter) && strings.HasPrefix(o.Model().ID, "anthropic/")
 70	// Add system message first
 71	systemMessage := o.providerOptions.systemMessage
 72	if o.providerOptions.systemPromptPrefix != "" {
 73		systemMessage = o.providerOptions.systemPromptPrefix + "\n" + systemMessage
 74	}
 75
 76	system := openai.SystemMessage(systemMessage)
 77	if isAnthropicModel && !o.providerOptions.disableCache {
 78		systemTextBlock := openai.ChatCompletionContentPartTextParam{Text: systemMessage}
 79		systemTextBlock.SetExtraFields(
 80			map[string]any{
 81				"cache_control": map[string]string{
 82					"type": "ephemeral",
 83				},
 84			},
 85		)
 86		var content []openai.ChatCompletionContentPartTextParam
 87		content = append(content, systemTextBlock)
 88		system = openai.SystemMessage(content)
 89	}
 90	openaiMessages = append(openaiMessages, system)
 91
 92	for i, msg := range messages {
 93		cache := false
 94		if i > len(messages)-3 {
 95			cache = true
 96		}
 97		switch msg.Role {
 98		case message.User:
 99			var content []openai.ChatCompletionContentPartUnionParam
100
101			textBlock := openai.ChatCompletionContentPartTextParam{Text: msg.Content().String()}
102			content = append(content, openai.ChatCompletionContentPartUnionParam{OfText: &textBlock})
103			hasBinaryContent := false
104			for _, binaryContent := range msg.BinaryContent() {
105				hasBinaryContent = true
106				imageURL := openai.ChatCompletionContentPartImageImageURLParam{URL: binaryContent.String(catwalk.InferenceProviderOpenAI)}
107				imageBlock := openai.ChatCompletionContentPartImageParam{ImageURL: imageURL}
108
109				content = append(content, openai.ChatCompletionContentPartUnionParam{OfImageURL: &imageBlock})
110			}
111			if cache && !o.providerOptions.disableCache && isAnthropicModel {
112				textBlock.SetExtraFields(map[string]any{
113					"cache_control": map[string]string{
114						"type": "ephemeral",
115					},
116				})
117			}
118			if hasBinaryContent || (isAnthropicModel && !o.providerOptions.disableCache) {
119				openaiMessages = append(openaiMessages, openai.UserMessage(content))
120			} else {
121				openaiMessages = append(openaiMessages, openai.UserMessage(msg.Content().String()))
122			}
123
124		case message.Assistant:
125			assistantMsg := openai.ChatCompletionAssistantMessageParam{
126				Role: "assistant",
127			}
128
129			hasContent := false
130			if msg.Content().String() != "" {
131				hasContent = true
132				textBlock := openai.ChatCompletionContentPartTextParam{Text: msg.Content().String()}
133				if cache && !o.providerOptions.disableCache && isAnthropicModel {
134					textBlock.SetExtraFields(map[string]any{
135						"cache_control": map[string]string{
136							"type": "ephemeral",
137						},
138					})
139				}
140				assistantMsg.Content = openai.ChatCompletionAssistantMessageParamContentUnion{
141					OfArrayOfContentParts: []openai.ChatCompletionAssistantMessageParamContentArrayOfContentPartUnion{
142						{
143							OfText: &textBlock,
144						},
145					},
146				}
147				if !isAnthropicModel {
148					assistantMsg.Content = openai.ChatCompletionAssistantMessageParamContentUnion{
149						OfString: param.NewOpt(msg.Content().String()),
150					}
151				}
152			}
153
154			if len(msg.ToolCalls()) > 0 {
155				hasContent = true
156				assistantMsg.ToolCalls = make([]openai.ChatCompletionMessageToolCallParam, len(msg.ToolCalls()))
157				for i, call := range msg.ToolCalls() {
158					assistantMsg.ToolCalls[i] = openai.ChatCompletionMessageToolCallParam{
159						ID:   call.ID,
160						Type: "function",
161						Function: openai.ChatCompletionMessageToolCallFunctionParam{
162							Name:      call.Name,
163							Arguments: call.Input,
164						},
165					}
166				}
167			}
168			if !hasContent {
169				slog.Warn("There is a message without content, investigate, this should not happen")
170				continue
171			}
172
173			openaiMessages = append(openaiMessages, openai.ChatCompletionMessageParamUnion{
174				OfAssistant: &assistantMsg,
175			})
176
177		case message.Tool:
178			for _, result := range msg.ToolResults() {
179				openaiMessages = append(openaiMessages,
180					openai.ToolMessage(result.Content, result.ToolCallID),
181				)
182			}
183		}
184	}
185
186	return
187}
188
189func (o *openaiClient) convertTools(tools []tools.BaseTool) []openai.ChatCompletionToolParam {
190	openaiTools := make([]openai.ChatCompletionToolParam, len(tools))
191
192	for i, tool := range tools {
193		info := tool.Info()
194		openaiTools[i] = openai.ChatCompletionToolParam{
195			Function: openai.FunctionDefinitionParam{
196				Name:        info.Name,
197				Description: openai.String(info.Description),
198				Parameters: openai.FunctionParameters{
199					"type":       "object",
200					"properties": info.Parameters,
201					"required":   info.Required,
202				},
203			},
204		}
205	}
206
207	return openaiTools
208}
209
210func (o *openaiClient) finishReason(reason string) message.FinishReason {
211	switch reason {
212	case "stop":
213		return message.FinishReasonEndTurn
214	case "length":
215		return message.FinishReasonMaxTokens
216	case "tool_calls":
217		return message.FinishReasonToolUse
218	default:
219		return message.FinishReasonUnknown
220	}
221}
222
223func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessageParamUnion, tools []openai.ChatCompletionToolParam) openai.ChatCompletionNewParams {
224	model := o.providerOptions.model(o.providerOptions.modelType)
225	cfg := config.Get()
226
227	modelConfig := cfg.Models[config.SelectedModelTypeLarge]
228	if o.providerOptions.modelType == config.SelectedModelTypeSmall {
229		modelConfig = cfg.Models[config.SelectedModelTypeSmall]
230	}
231
232	reasoningEffort := modelConfig.ReasoningEffort
233
234	params := openai.ChatCompletionNewParams{
235		Model:    openai.ChatModel(model.ID),
236		Messages: messages,
237		Tools:    tools,
238	}
239
240	maxTokens := model.DefaultMaxTokens
241	if modelConfig.MaxTokens > 0 {
242		maxTokens = modelConfig.MaxTokens
243	}
244
245	// Override max tokens if set in provider options
246	if o.providerOptions.maxTokens > 0 {
247		maxTokens = o.providerOptions.maxTokens
248	}
249	if model.CanReason {
250		params.MaxCompletionTokens = openai.Int(maxTokens)
251		switch reasoningEffort {
252		case "low":
253			params.ReasoningEffort = shared.ReasoningEffortLow
254		case "medium":
255			params.ReasoningEffort = shared.ReasoningEffortMedium
256		case "high":
257			params.ReasoningEffort = shared.ReasoningEffortHigh
258		case "minimal":
259			params.ReasoningEffort = shared.ReasoningEffort("minimal")
260		default:
261			params.ReasoningEffort = shared.ReasoningEffort(reasoningEffort)
262		}
263	} else {
264		params.MaxTokens = openai.Int(maxTokens)
265	}
266
267	return params
268}
269
270func (o *openaiClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (response *ProviderResponse, err error) {
271	params := o.preparedParams(o.convertMessages(messages), o.convertTools(tools))
272	attempts := 0
273	for {
274		attempts++
275		openaiResponse, err := o.client.Chat.Completions.New(
276			ctx,
277			params,
278		)
279		// If there is an error we are going to see if we can retry the call
280		if err != nil {
281			retry, after, retryErr := o.shouldRetry(attempts, err)
282			if retryErr != nil {
283				return nil, retryErr
284			}
285			if retry {
286				slog.Warn("Retrying due to rate limit", "attempt", attempts, "max_retries", maxRetries)
287				select {
288				case <-ctx.Done():
289					return nil, ctx.Err()
290				case <-time.After(time.Duration(after) * time.Millisecond):
291					continue
292				}
293			}
294			return nil, retryErr
295		}
296
297		if len(openaiResponse.Choices) == 0 {
298			return nil, fmt.Errorf("received empty response from OpenAI API - check endpoint configuration")
299		}
300
301		content := ""
302		if openaiResponse.Choices[0].Message.Content != "" {
303			content = openaiResponse.Choices[0].Message.Content
304		}
305
306		toolCalls := o.toolCalls(*openaiResponse)
307		finishReason := o.finishReason(string(openaiResponse.Choices[0].FinishReason))
308
309		if len(toolCalls) > 0 {
310			finishReason = message.FinishReasonToolUse
311		}
312
313		return &ProviderResponse{
314			Content:      content,
315			ToolCalls:    toolCalls,
316			Usage:        o.usage(*openaiResponse),
317			FinishReason: finishReason,
318		}, nil
319	}
320}
321
322func (o *openaiClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {
323	params := o.preparedParams(o.convertMessages(messages), o.convertTools(tools))
324	params.StreamOptions = openai.ChatCompletionStreamOptionsParam{
325		IncludeUsage: openai.Bool(true),
326	}
327
328	attempts := 0
329	eventChan := make(chan ProviderEvent)
330
331	go func() {
332		for {
333			attempts++
334			// Kujtim: fixes an issue with anthropig models on openrouter
335			if len(params.Tools) == 0 {
336				params.Tools = nil
337			}
338			openaiStream := o.client.Chat.Completions.NewStreaming(
339				ctx,
340				params,
341			)
342
343			acc := openai.ChatCompletionAccumulator{}
344			currentContent := ""
345			toolCalls := make([]message.ToolCall, 0)
346			var msgToolCalls []openai.ChatCompletionMessageToolCall
347			for openaiStream.Next() {
348				chunk := openaiStream.Current()
349				// Kujtim: this is an issue with openrouter qwen, its sending -1 for the tool index
350				if len(chunk.Choices) > 0 && len(chunk.Choices[0].Delta.ToolCalls) > 0 && chunk.Choices[0].Delta.ToolCalls[0].Index == -1 {
351					chunk.Choices[0].Delta.ToolCalls[0].Index = 0
352				}
353				acc.AddChunk(chunk)
354				for i, choice := range chunk.Choices {
355					reasoning, ok := choice.Delta.JSON.ExtraFields["reasoning"]
356					if ok && reasoning.Raw() != "" {
357						reasoningStr := ""
358						json.Unmarshal([]byte(reasoning.Raw()), &reasoningStr)
359						if reasoningStr != "" {
360							eventChan <- ProviderEvent{
361								Type:     EventThinkingDelta,
362								Thinking: reasoningStr,
363							}
364						}
365					}
366					if choice.Delta.Content != "" {
367						eventChan <- ProviderEvent{
368							Type:    EventContentDelta,
369							Content: choice.Delta.Content,
370						}
371						currentContent += choice.Delta.Content
372					} else if len(choice.Delta.ToolCalls) > 0 {
373						toolCall := choice.Delta.ToolCalls[0]
374						newToolCall := false
375						if len(msgToolCalls)-1 >= int(toolCall.Index) { // tool call exists
376							existingToolCall := msgToolCalls[toolCall.Index]
377							if toolCall.ID != "" && toolCall.ID != existingToolCall.ID {
378								found := false
379								// try to find the tool based on the ID
380								for i, tool := range msgToolCalls {
381									if tool.ID == toolCall.ID {
382										msgToolCalls[i].Function.Arguments += toolCall.Function.Arguments
383										found = true
384									}
385								}
386								if !found {
387									newToolCall = true
388								}
389							} else {
390								msgToolCalls[toolCall.Index].Function.Arguments += toolCall.Function.Arguments
391							}
392						} else {
393							newToolCall = true
394						}
395						if newToolCall { // new tool call
396							if toolCall.ID == "" {
397								toolCall.ID = uuid.NewString()
398							}
399							eventChan <- ProviderEvent{
400								Type: EventToolUseStart,
401								ToolCall: &message.ToolCall{
402									ID:       toolCall.ID,
403									Name:     toolCall.Function.Name,
404									Finished: false,
405								},
406							}
407							msgToolCalls = append(msgToolCalls, openai.ChatCompletionMessageToolCall{
408								ID:   toolCall.ID,
409								Type: "function",
410								Function: openai.ChatCompletionMessageToolCallFunction{
411									Name:      toolCall.Function.Name,
412									Arguments: toolCall.Function.Arguments,
413								},
414							})
415						}
416					}
417					acc.Choices[i].Message.ToolCalls = slices.Clone(msgToolCalls)
418				}
419			}
420
421			err := openaiStream.Err()
422			if err == nil || errors.Is(err, io.EOF) {
423				if len(acc.Choices) == 0 {
424					eventChan <- ProviderEvent{
425						Type:  EventError,
426						Error: fmt.Errorf("received empty streaming response from OpenAI API - check endpoint configuration"),
427					}
428					return
429				}
430
431				resultFinishReason := acc.Choices[0].FinishReason
432				if resultFinishReason == "" {
433					// If the finish reason is empty, we assume it was a successful completion
434					// INFO: this is happening for openrouter for some reason
435					resultFinishReason = "stop"
436				}
437				// Stream completed successfully
438				finishReason := o.finishReason(resultFinishReason)
439				if len(acc.Choices[0].Message.ToolCalls) > 0 {
440					toolCalls = append(toolCalls, o.toolCalls(acc.ChatCompletion)...)
441				}
442				if len(toolCalls) > 0 {
443					finishReason = message.FinishReasonToolUse
444				}
445
446				eventChan <- ProviderEvent{
447					Type: EventComplete,
448					Response: &ProviderResponse{
449						Content:      currentContent,
450						ToolCalls:    toolCalls,
451						Usage:        o.usage(acc.ChatCompletion),
452						FinishReason: finishReason,
453					},
454				}
455				close(eventChan)
456				return
457			}
458
459			// If there is an error we are going to see if we can retry the call
460			retry, after, retryErr := o.shouldRetry(attempts, err)
461			if retryErr != nil {
462				eventChan <- ProviderEvent{Type: EventError, Error: retryErr}
463				close(eventChan)
464				return
465			}
466			if retry {
467				slog.Warn("Retrying due to rate limit", "attempt", attempts, "max_retries", maxRetries)
468				select {
469				case <-ctx.Done():
470					// context cancelled
471					if ctx.Err() == nil {
472						eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
473					}
474					close(eventChan)
475					return
476				case <-time.After(time.Duration(after) * time.Millisecond):
477					continue
478				}
479			}
480			eventChan <- ProviderEvent{Type: EventError, Error: retryErr}
481			close(eventChan)
482			return
483		}
484	}()
485
486	return eventChan
487}
488
489func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error) {
490	if attempts > maxRetries {
491		return false, 0, fmt.Errorf("maximum retry attempts reached for rate limit: %d retries", maxRetries)
492	}
493	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
494		return false, 0, err
495	}
496	var apiErr *openai.Error
497	retryMs := 0
498	retryAfterValues := []string{}
499	if errors.As(err, &apiErr) {
500		// Check for token expiration (401 Unauthorized)
501		if apiErr.StatusCode == 401 {
502			o.providerOptions.apiKey, err = config.Get().Resolve(o.providerOptions.config.APIKey)
503			if err != nil {
504				return false, 0, fmt.Errorf("failed to resolve API key: %w", err)
505			}
506			o.client = createOpenAIClient(o.providerOptions)
507			return true, 0, nil
508		}
509
510		if apiErr.StatusCode != 429 && apiErr.StatusCode != 500 {
511			return false, 0, err
512		}
513
514		retryAfterValues = apiErr.Response.Header.Values("Retry-After")
515	}
516
517	if apiErr != nil {
518		slog.Warn("OpenAI API error", "status_code", apiErr.StatusCode, "message", apiErr.Message, "type", apiErr.Type)
519		if len(retryAfterValues) > 0 {
520			slog.Warn("Retry-After header", "values", retryAfterValues)
521		}
522	} else {
523		slog.Error("OpenAI API error", "error", err.Error(), "attempt", attempts, "max_retries", maxRetries)
524	}
525
526	backoffMs := 2000 * (1 << (attempts - 1))
527	jitterMs := int(float64(backoffMs) * 0.2)
528	retryMs = backoffMs + jitterMs
529	if len(retryAfterValues) > 0 {
530		if _, err := fmt.Sscanf(retryAfterValues[0], "%d", &retryMs); err == nil {
531			retryMs = retryMs * 1000
532		}
533	}
534	return true, int64(retryMs), nil
535}
536
537func (o *openaiClient) toolCalls(completion openai.ChatCompletion) []message.ToolCall {
538	var toolCalls []message.ToolCall
539
540	if len(completion.Choices) > 0 && len(completion.Choices[0].Message.ToolCalls) > 0 {
541		for _, call := range completion.Choices[0].Message.ToolCalls {
542			toolCall := message.ToolCall{
543				ID:       call.ID,
544				Name:     call.Function.Name,
545				Input:    call.Function.Arguments,
546				Type:     "function",
547				Finished: true,
548			}
549			toolCalls = append(toolCalls, toolCall)
550		}
551	}
552
553	return toolCalls
554}
555
556func (o *openaiClient) usage(completion openai.ChatCompletion) TokenUsage {
557	cachedTokens := completion.Usage.PromptTokensDetails.CachedTokens
558	inputTokens := completion.Usage.PromptTokens - cachedTokens
559
560	return TokenUsage{
561		InputTokens:         inputTokens,
562		OutputTokens:        completion.Usage.CompletionTokens,
563		CacheCreationTokens: 0, // OpenAI doesn't provide this directly
564		CacheReadTokens:     cachedTokens,
565	}
566}
567
568func (o *openaiClient) Model() catwalk.Model {
569	return o.providerOptions.model(o.providerOptions.modelType)
570}