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