1package provider
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "io"
9 "time"
10
11 "github.com/charmbracelet/crush/internal/config"
12 "github.com/charmbracelet/crush/internal/fur/provider"
13 "github.com/charmbracelet/crush/internal/llm/tools"
14 "github.com/charmbracelet/crush/internal/logging"
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 openaiClientOptions = append(openaiClientOptions, option.WithBaseURL(opts.baseURL))
42 }
43
44 if opts.extraHeaders != nil {
45 for key, value := range opts.extraHeaders {
46 openaiClientOptions = append(openaiClientOptions, option.WithHeader(key, value))
47 }
48 }
49
50 return openai.NewClient(openaiClientOptions...)
51}
52
53func (o *openaiClient) convertMessages(messages []message.Message) (openaiMessages []openai.ChatCompletionMessageParamUnion) {
54 // Add system message first
55 openaiMessages = append(openaiMessages, openai.SystemMessage(o.providerOptions.systemMessage))
56
57 for _, msg := range messages {
58 switch msg.Role {
59 case message.User:
60 var content []openai.ChatCompletionContentPartUnionParam
61 textBlock := openai.ChatCompletionContentPartTextParam{Text: msg.Content().String()}
62 content = append(content, openai.ChatCompletionContentPartUnionParam{OfText: &textBlock})
63 for _, binaryContent := range msg.BinaryContent() {
64 imageURL := openai.ChatCompletionContentPartImageImageURLParam{URL: binaryContent.String(provider.InferenceProviderOpenAI)}
65 imageBlock := openai.ChatCompletionContentPartImageParam{ImageURL: imageURL}
66
67 content = append(content, openai.ChatCompletionContentPartUnionParam{OfImageURL: &imageBlock})
68 }
69
70 openaiMessages = append(openaiMessages, openai.UserMessage(content))
71
72 case message.Assistant:
73 assistantMsg := openai.ChatCompletionAssistantMessageParam{
74 Role: "assistant",
75 }
76
77 if msg.Content().String() != "" {
78 assistantMsg.Content = openai.ChatCompletionAssistantMessageParamContentUnion{
79 OfString: openai.String(msg.Content().String()),
80 }
81 }
82
83 if len(msg.ToolCalls()) > 0 {
84 assistantMsg.ToolCalls = make([]openai.ChatCompletionMessageToolCallParam, len(msg.ToolCalls()))
85 for i, call := range msg.ToolCalls() {
86 assistantMsg.ToolCalls[i] = openai.ChatCompletionMessageToolCallParam{
87 ID: call.ID,
88 Type: "function",
89 Function: openai.ChatCompletionMessageToolCallFunctionParam{
90 Name: call.Name,
91 Arguments: call.Input,
92 },
93 }
94 }
95 }
96
97 openaiMessages = append(openaiMessages, openai.ChatCompletionMessageParamUnion{
98 OfAssistant: &assistantMsg,
99 })
100
101 case message.Tool:
102 for _, result := range msg.ToolResults() {
103 openaiMessages = append(openaiMessages,
104 openai.ToolMessage(result.Content, result.ToolCallID),
105 )
106 }
107 }
108 }
109
110 return
111}
112
113func (o *openaiClient) convertTools(tools []tools.BaseTool) []openai.ChatCompletionToolParam {
114 openaiTools := make([]openai.ChatCompletionToolParam, len(tools))
115
116 for i, tool := range tools {
117 info := tool.Info()
118 openaiTools[i] = openai.ChatCompletionToolParam{
119 Function: openai.FunctionDefinitionParam{
120 Name: info.Name,
121 Description: openai.String(info.Description),
122 Parameters: openai.FunctionParameters{
123 "type": "object",
124 "properties": info.Parameters,
125 "required": info.Required,
126 },
127 },
128 }
129 }
130
131 return openaiTools
132}
133
134func (o *openaiClient) finishReason(reason string) message.FinishReason {
135 switch reason {
136 case "stop":
137 return message.FinishReasonEndTurn
138 case "length":
139 return message.FinishReasonMaxTokens
140 case "tool_calls":
141 return message.FinishReasonToolUse
142 default:
143 return message.FinishReasonUnknown
144 }
145}
146
147func (o *openaiClient) preparedParams(messages []openai.ChatCompletionMessageParamUnion, tools []openai.ChatCompletionToolParam) openai.ChatCompletionNewParams {
148 model := o.providerOptions.model(o.providerOptions.modelType)
149 cfg := config.Get()
150
151 modelConfig := cfg.Models.Large
152 if o.providerOptions.modelType == config.SmallModel {
153 modelConfig = cfg.Models.Small
154 }
155
156 reasoningEffort := model.ReasoningEffort
157 if modelConfig.ReasoningEffort != "" {
158 reasoningEffort = modelConfig.ReasoningEffort
159 }
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 logging.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 logging.WarnPersist(fmt.Sprintf("Retrying due to rate limit... attempt %d of %d", attempts, maxRetries), logging.PersistTimeArg, time.Millisecond*time.Duration(after+100))
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 logging.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 // Stream completed successfully
293 finishReason := o.finishReason(string(acc.ChatCompletion.Choices[0].FinishReason))
294 if len(acc.Choices[0].Message.ToolCalls) > 0 {
295 toolCalls = append(toolCalls, o.toolCalls(acc.ChatCompletion)...)
296 }
297 if len(toolCalls) > 0 {
298 finishReason = message.FinishReasonToolUse
299 }
300
301 eventChan <- ProviderEvent{
302 Type: EventComplete,
303 Response: &ProviderResponse{
304 Content: currentContent,
305 ToolCalls: toolCalls,
306 Usage: o.usage(acc.ChatCompletion),
307 FinishReason: finishReason,
308 },
309 }
310 close(eventChan)
311 return
312 }
313
314 // If there is an error we are going to see if we can retry the call
315 retry, after, retryErr := o.shouldRetry(attempts, err)
316 if retryErr != nil {
317 eventChan <- ProviderEvent{Type: EventError, Error: retryErr}
318 close(eventChan)
319 return
320 }
321 if retry {
322 logging.WarnPersist(fmt.Sprintf("Retrying due to rate limit... attempt %d of %d", attempts, maxRetries), logging.PersistTimeArg, time.Millisecond*time.Duration(after+100))
323 select {
324 case <-ctx.Done():
325 // context cancelled
326 if ctx.Err() == nil {
327 eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
328 }
329 close(eventChan)
330 return
331 case <-time.After(time.Duration(after) * time.Millisecond):
332 continue
333 }
334 }
335 eventChan <- ProviderEvent{Type: EventError, Error: retryErr}
336 close(eventChan)
337 return
338 }
339 }()
340
341 return eventChan
342}
343
344func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error) {
345 var apiErr *openai.Error
346 if !errors.As(err, &apiErr) {
347 return false, 0, err
348 }
349
350 // Check for token expiration (401 Unauthorized)
351 if apiErr.StatusCode == 401 {
352 o.providerOptions.apiKey, err = config.ResolveAPIKey(o.providerOptions.config.APIKey)
353 if err != nil {
354 return false, 0, fmt.Errorf("failed to resolve API key: %w", err)
355 }
356 o.client = createOpenAIClient(o.providerOptions)
357 return true, 0, nil
358 }
359
360 if apiErr.StatusCode != 429 && apiErr.StatusCode != 500 {
361 return false, 0, err
362 }
363
364 if attempts > maxRetries {
365 return false, 0, fmt.Errorf("maximum retry attempts reached for rate limit: %d retries", maxRetries)
366 }
367
368 retryMs := 0
369 retryAfterValues := apiErr.Response.Header.Values("Retry-After")
370
371 backoffMs := 2000 * (1 << (attempts - 1))
372 jitterMs := int(float64(backoffMs) * 0.2)
373 retryMs = backoffMs + jitterMs
374 if len(retryAfterValues) > 0 {
375 if _, err := fmt.Sscanf(retryAfterValues[0], "%d", &retryMs); err == nil {
376 retryMs = retryMs * 1000
377 }
378 }
379 return true, int64(retryMs), nil
380}
381
382func (o *openaiClient) toolCalls(completion openai.ChatCompletion) []message.ToolCall {
383 var toolCalls []message.ToolCall
384
385 if len(completion.Choices) > 0 && len(completion.Choices[0].Message.ToolCalls) > 0 {
386 for _, call := range completion.Choices[0].Message.ToolCalls {
387 toolCall := message.ToolCall{
388 ID: call.ID,
389 Name: call.Function.Name,
390 Input: call.Function.Arguments,
391 Type: "function",
392 Finished: true,
393 }
394 toolCalls = append(toolCalls, toolCall)
395 }
396 }
397
398 return toolCalls
399}
400
401func (o *openaiClient) usage(completion openai.ChatCompletion) TokenUsage {
402 cachedTokens := completion.Usage.PromptTokensDetails.CachedTokens
403 inputTokens := completion.Usage.PromptTokens - cachedTokens
404
405 return TokenUsage{
406 InputTokens: inputTokens,
407 OutputTokens: completion.Usage.CompletionTokens,
408 CacheCreationTokens: 0, // OpenAI doesn't provide this directly
409 CacheReadTokens: cachedTokens,
410 }
411}
412
413func (a *openaiClient) Model() config.Model {
414 return a.providerOptions.model(a.providerOptions.modelType)
415}