openai.go 23.7 KB
Newer Older
1
2
3
4
5
// openai package provides middleware for partial compatibility with the OpenAI REST API
package openai

import (
	"bytes"
6
	"encoding/base64"
7
	"encoding/json"
Michael Yang's avatar
lint  
Michael Yang committed
8
	"errors"
9
10
	"fmt"
	"io"
royjhan's avatar
royjhan committed
11
	"log/slog"
12
13
	"math/rand"
	"net/http"
14
	"strings"
15
16
17
	"time"

	"github.com/gin-gonic/gin"
Michael Yang's avatar
lint  
Michael Yang committed
18

19
	"github.com/ollama/ollama/api"
20
	"github.com/ollama/ollama/types/model"
21
22
23
24
25
26
27
28
29
30
31
32
33
34
)

type Error struct {
	Message string      `json:"message"`
	Type    string      `json:"type"`
	Param   interface{} `json:"param"`
	Code    *string     `json:"code"`
}

type ErrorResponse struct {
	Error Error `json:"error"`
}

type Message struct {
royjhan's avatar
royjhan committed
35
36
37
	Role      string     `json:"role"`
	Content   any        `json:"content"`
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
38
39
40
41
42
43
44
45
46
47
48
49
50
51
}

type Choice struct {
	Index        int     `json:"index"`
	Message      Message `json:"message"`
	FinishReason *string `json:"finish_reason"`
}

type ChunkChoice struct {
	Index        int     `json:"index"`
	Delta        Message `json:"delta"`
	FinishReason *string `json:"finish_reason"`
}

52
53
54
55
56
57
type CompleteChunkChoice struct {
	Text         string  `json:"text"`
	Index        int     `json:"index"`
	FinishReason *string `json:"finish_reason"`
}

58
59
60
61
62
63
64
type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type ResponseFormat struct {
65
66
67
68
69
	Type       string      `json:"type"`
	JsonSchema *JsonSchema `json:"json_schema,omitempty"`
}

type JsonSchema struct {
70
	Schema json.RawMessage `json:"schema"`
71
72
}

73
74
75
76
77
type EmbedRequest struct {
	Input any    `json:"input"`
	Model string `json:"model"`
}

78
79
80
81
type StreamOptions struct {
	IncludeUsage bool `json:"include_usage"`
}

82
83
84
85
type ChatCompletionRequest struct {
	Model            string          `json:"model"`
	Messages         []Message       `json:"messages"`
	Stream           bool            `json:"stream"`
86
	StreamOptions    *StreamOptions  `json:"stream_options"`
87
88
89
90
91
	MaxTokens        *int            `json:"max_tokens"`
	Seed             *int            `json:"seed"`
	Stop             any             `json:"stop"`
	Temperature      *float64        `json:"temperature"`
	FrequencyPenalty *float64        `json:"frequency_penalty"`
92
	PresencePenalty  *float64        `json:"presence_penalty"`
93
94
	TopP             *float64        `json:"top_p"`
	ResponseFormat   *ResponseFormat `json:"response_format"`
royjhan's avatar
royjhan committed
95
	Tools            []api.Tool      `json:"tools"`
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
}

type ChatCompletion struct {
	Id                string   `json:"id"`
	Object            string   `json:"object"`
	Created           int64    `json:"created"`
	Model             string   `json:"model"`
	SystemFingerprint string   `json:"system_fingerprint"`
	Choices           []Choice `json:"choices"`
	Usage             Usage    `json:"usage,omitempty"`
}

type ChatCompletionChunk struct {
	Id                string        `json:"id"`
	Object            string        `json:"object"`
	Created           int64         `json:"created"`
	Model             string        `json:"model"`
	SystemFingerprint string        `json:"system_fingerprint"`
	Choices           []ChunkChoice `json:"choices"`
115
	Usage             *Usage        `json:"usage,omitempty"`
116
117
}

118
119
// TODO (https://github.com/ollama/ollama/issues/5259): support []string, []int and [][]int
type CompletionRequest struct {
120
121
122
123
124
125
126
127
128
129
130
131
	Model            string         `json:"model"`
	Prompt           string         `json:"prompt"`
	FrequencyPenalty float32        `json:"frequency_penalty"`
	MaxTokens        *int           `json:"max_tokens"`
	PresencePenalty  float32        `json:"presence_penalty"`
	Seed             *int           `json:"seed"`
	Stop             any            `json:"stop"`
	Stream           bool           `json:"stream"`
	StreamOptions    *StreamOptions `json:"stream_options"`
	Temperature      *float32       `json:"temperature"`
	TopP             float32        `json:"top_p"`
	Suffix           string         `json:"suffix"`
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
}

type Completion struct {
	Id                string                `json:"id"`
	Object            string                `json:"object"`
	Created           int64                 `json:"created"`
	Model             string                `json:"model"`
	SystemFingerprint string                `json:"system_fingerprint"`
	Choices           []CompleteChunkChoice `json:"choices"`
	Usage             Usage                 `json:"usage,omitempty"`
}

type CompletionChunk struct {
	Id                string                `json:"id"`
	Object            string                `json:"object"`
	Created           int64                 `json:"created"`
	Choices           []CompleteChunkChoice `json:"choices"`
	Model             string                `json:"model"`
	SystemFingerprint string                `json:"system_fingerprint"`
151
	Usage             *Usage                `json:"usage,omitempty"`
152
153
}

royjhan's avatar
royjhan committed
154
155
type ToolCall struct {
	ID       string `json:"id"`
156
	Index    int    `json:"index"`
royjhan's avatar
royjhan committed
157
158
159
160
161
162
163
	Type     string `json:"type"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `json:"function"`
}

164
165
166
167
168
169
170
type Model struct {
	Id      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	OwnedBy string `json:"owned_by"`
}

171
172
173
174
175
176
type Embedding struct {
	Object    string    `json:"object"`
	Embedding []float32 `json:"embedding"`
	Index     int       `json:"index"`
}

177
178
179
180
181
type ListCompletion struct {
	Object string  `json:"object"`
	Data   []Model `json:"data"`
}

182
type EmbeddingList struct {
183
184
185
186
187
188
189
190
191
	Object string         `json:"object"`
	Data   []Embedding    `json:"data"`
	Model  string         `json:"model"`
	Usage  EmbeddingUsage `json:"usage,omitempty"`
}

type EmbeddingUsage struct {
	PromptTokens int `json:"prompt_tokens"`
	TotalTokens  int `json:"total_tokens"`
192
193
}

194
195
196
197
198
199
200
201
202
203
204
205
206
207
func NewError(code int, message string) ErrorResponse {
	var etype string
	switch code {
	case http.StatusBadRequest:
		etype = "invalid_request_error"
	case http.StatusNotFound:
		etype = "not_found_error"
	default:
		etype = "api_error"
	}

	return ErrorResponse{Error{Type: etype, Message: message}}
}

208
209
210
211
212
213
214
215
func toUsage(r api.ChatResponse) Usage {
	return Usage{
		PromptTokens:     r.PromptEvalCount,
		CompletionTokens: r.EvalCount,
		TotalTokens:      r.PromptEvalCount + r.EvalCount,
	}
}

royjhan's avatar
royjhan committed
216
217
218
219
220
221
222
223
224
func toolCallId() string {
	const letterBytes = "abcdefghijklmnopqrstuvwxyz0123456789"
	b := make([]byte, 8)
	for i := range b {
		b[i] = letterBytes[rand.Intn(len(letterBytes))]
	}
	return "call_" + strings.ToLower(string(b))
}

225
226
227
func toToolCalls(tc []api.ToolCall) []ToolCall {
	toolCalls := make([]ToolCall, len(tc))
	for i, tc := range tc {
royjhan's avatar
royjhan committed
228
229
230
		toolCalls[i].ID = toolCallId()
		toolCalls[i].Type = "function"
		toolCalls[i].Function.Name = tc.Function.Name
231
		toolCalls[i].Index = tc.Function.Index
royjhan's avatar
royjhan committed
232
233
234
235
236
237
238
239
240

		args, err := json.Marshal(tc.Function.Arguments)
		if err != nil {
			slog.Error("could not marshall function arguments to json", "error", err)
			continue
		}

		toolCalls[i].Function.Arguments = string(args)
	}
241
242
	return toolCalls
}
royjhan's avatar
royjhan committed
243

244
245
func toChatCompletion(id string, r api.ChatResponse) ChatCompletion {
	toolCalls := toToolCalls(r.Message.ToolCalls)
246
247
248
249
250
251
252
	return ChatCompletion{
		Id:                id,
		Object:            "chat.completion",
		Created:           r.CreatedAt.Unix(),
		Model:             r.Model,
		SystemFingerprint: "fp_ollama",
		Choices: []Choice{{
253
			Index:   0,
royjhan's avatar
royjhan committed
254
			Message: Message{Role: r.Message.Role, Content: r.Message.Content, ToolCalls: toolCalls},
255
			FinishReason: func(reason string) *string {
256
257
258
				if len(toolCalls) > 0 {
					reason = "tool_calls"
				}
259
260
261
262
263
				if len(reason) > 0 {
					return &reason
				}
				return nil
			}(r.DoneReason),
264
		}},
265
		Usage: toUsage(r),
266
267
268
269
	}
}

func toChunk(id string, r api.ChatResponse) ChatCompletionChunk {
270
	toolCalls := toToolCalls(r.Message.ToolCalls)
271
272
273
274
275
276
	return ChatCompletionChunk{
		Id:                id,
		Object:            "chat.completion.chunk",
		Created:           time.Now().Unix(),
		Model:             r.Model,
		SystemFingerprint: "fp_ollama",
277
278
		Choices: []ChunkChoice{{
			Index: 0,
279
			Delta: Message{Role: "assistant", Content: r.Message.Content, ToolCalls: toolCalls},
280
281
282
283
284
285
286
			FinishReason: func(reason string) *string {
				if len(reason) > 0 {
					return &reason
				}
				return nil
			}(r.DoneReason),
		}},
287
288
289
	}
}

290
291
292
293
294
295
296
297
func toUsageGenerate(r api.GenerateResponse) Usage {
	return Usage{
		PromptTokens:     r.PromptEvalCount,
		CompletionTokens: r.EvalCount,
		TotalTokens:      r.PromptEvalCount + r.EvalCount,
	}
}

298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
func toCompletion(id string, r api.GenerateResponse) Completion {
	return Completion{
		Id:                id,
		Object:            "text_completion",
		Created:           r.CreatedAt.Unix(),
		Model:             r.Model,
		SystemFingerprint: "fp_ollama",
		Choices: []CompleteChunkChoice{{
			Text:  r.Response,
			Index: 0,
			FinishReason: func(reason string) *string {
				if len(reason) > 0 {
					return &reason
				}
				return nil
			}(r.DoneReason),
		}},
315
		Usage: toUsageGenerate(r),
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
	}
}

func toCompleteChunk(id string, r api.GenerateResponse) CompletionChunk {
	return CompletionChunk{
		Id:                id,
		Object:            "text_completion",
		Created:           time.Now().Unix(),
		Model:             r.Model,
		SystemFingerprint: "fp_ollama",
		Choices: []CompleteChunkChoice{{
			Text:  r.Response,
			Index: 0,
			FinishReason: func(reason string) *string {
				if len(reason) > 0 {
					return &reason
				}
				return nil
			}(r.DoneReason),
		}},
	}
}

339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
func toListCompletion(r api.ListResponse) ListCompletion {
	var data []Model
	for _, m := range r.Models {
		data = append(data, Model{
			Id:      m.Name,
			Object:  "model",
			Created: m.ModifiedAt.Unix(),
			OwnedBy: model.ParseName(m.Name).Namespace,
		})
	}

	return ListCompletion{
		Object: "list",
		Data:   data,
	}
}

356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
func toEmbeddingList(model string, r api.EmbedResponse) EmbeddingList {
	if r.Embeddings != nil {
		var data []Embedding
		for i, e := range r.Embeddings {
			data = append(data, Embedding{
				Object:    "embedding",
				Embedding: e,
				Index:     i,
			})
		}

		return EmbeddingList{
			Object: "list",
			Data:   data,
			Model:  model,
371
372
373
374
			Usage: EmbeddingUsage{
				PromptTokens: r.PromptEvalCount,
				TotalTokens:  r.PromptEvalCount,
			},
375
376
377
378
379
380
		}
	}

	return EmbeddingList{}
}

381
382
383
384
385
386
387
388
389
func toModel(r api.ShowResponse, m string) Model {
	return Model{
		Id:      m,
		Object:  "model",
		Created: r.ModifiedAt.Unix(),
		OwnedBy: model.ParseName(m).Namespace,
	}
}

390
func fromChatRequest(r ChatCompletionRequest) (*api.ChatRequest, error) {
391
392
	var messages []api.Message
	for _, msg := range r.Messages {
393
394
395
396
397
398
399
		switch content := msg.Content.(type) {
		case string:
			messages = append(messages, api.Message{Role: msg.Role, Content: content})
		case []any:
			for _, c := range content {
				data, ok := c.(map[string]any)
				if !ok {
Michael Yang's avatar
lint  
Michael Yang committed
400
					return nil, errors.New("invalid message format")
401
402
403
404
405
				}
				switch data["type"] {
				case "text":
					text, ok := data["text"].(string)
					if !ok {
Michael Yang's avatar
lint  
Michael Yang committed
406
						return nil, errors.New("invalid message format")
407
					}
408
					messages = append(messages, api.Message{Role: msg.Role, Content: text})
409
410
411
412
				case "image_url":
					var url string
					if urlMap, ok := data["image_url"].(map[string]any); ok {
						if url, ok = urlMap["url"].(string); !ok {
Michael Yang's avatar
lint  
Michael Yang committed
413
							return nil, errors.New("invalid message format")
414
415
416
						}
					} else {
						if url, ok = data["image_url"].(string); !ok {
Michael Yang's avatar
lint  
Michael Yang committed
417
							return nil, errors.New("invalid message format")
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
						}
					}

					types := []string{"jpeg", "jpg", "png"}
					valid := false
					for _, t := range types {
						prefix := "data:image/" + t + ";base64,"
						if strings.HasPrefix(url, prefix) {
							url = strings.TrimPrefix(url, prefix)
							valid = true
							break
						}
					}

					if !valid {
Michael Yang's avatar
lint  
Michael Yang committed
433
						return nil, errors.New("invalid image input")
434
435
436
437
					}

					img, err := base64.StdEncoding.DecodeString(url)
					if err != nil {
Michael Yang's avatar
lint  
Michael Yang committed
438
						return nil, errors.New("invalid message format")
439
					}
440
441

					messages = append(messages, api.Message{Role: msg.Role, Images: []api.ImageData{img}})
442
				default:
Michael Yang's avatar
lint  
Michael Yang committed
443
					return nil, errors.New("invalid message format")
444
445
446
				}
			}
		default:
royjhan's avatar
royjhan committed
447
448
449
450
451
452
453
454
455
			if msg.ToolCalls == nil {
				return nil, fmt.Errorf("invalid message content type: %T", content)
			}

			toolCalls := make([]api.ToolCall, len(msg.ToolCalls))
			for i, tc := range msg.ToolCalls {
				toolCalls[i].Function.Name = tc.Function.Name
				err := json.Unmarshal([]byte(tc.Function.Arguments), &toolCalls[i].Function.Arguments)
				if err != nil {
Michael Yang's avatar
lint  
Michael Yang committed
456
					return nil, errors.New("invalid tool call arguments")
royjhan's avatar
royjhan committed
457
458
459
				}
			}
			messages = append(messages, api.Message{Role: msg.Role, ToolCalls: toolCalls})
460
		}
461
462
463
464
465
466
467
	}

	options := make(map[string]interface{})

	switch stop := r.Stop.(type) {
	case string:
		options["stop"] = []string{stop}
468
	case []any:
469
470
471
472
473
474
475
476
477
478
479
480
481
482
		var stops []string
		for _, s := range stop {
			if str, ok := s.(string); ok {
				stops = append(stops, str)
			}
		}
		options["stop"] = stops
	}

	if r.MaxTokens != nil {
		options["num_predict"] = *r.MaxTokens
	}

	if r.Temperature != nil {
483
		options["temperature"] = *r.Temperature
484
485
486
487
488
489
490
491
492
	} else {
		options["temperature"] = 1.0
	}

	if r.Seed != nil {
		options["seed"] = *r.Seed
	}

	if r.FrequencyPenalty != nil {
493
		options["frequency_penalty"] = *r.FrequencyPenalty
494
495
496
	}

	if r.PresencePenalty != nil {
497
		options["presence_penalty"] = *r.PresencePenalty
498
499
500
501
502
503
504
505
	}

	if r.TopP != nil {
		options["top_p"] = *r.TopP
	} else {
		options["top_p"] = 1.0
	}

506
507
508
509
510
511
512
513
	var format json.RawMessage
	if r.ResponseFormat != nil {
		switch strings.ToLower(strings.TrimSpace(r.ResponseFormat.Type)) {
		// Support the old "json_object" type for OpenAI compatibility
		case "json_object":
			format = json.RawMessage(`"json"`)
		case "json_schema":
			if r.ResponseFormat.JsonSchema != nil {
514
				format = r.ResponseFormat.JsonSchema.Schema
515
516
			}
		}
517
518
	}

519
	return &api.ChatRequest{
520
521
522
523
524
		Model:    r.Model,
		Messages: messages,
		Format:   format,
		Options:  options,
		Stream:   &r.Stream,
royjhan's avatar
royjhan committed
525
		Tools:    r.Tools,
526
	}, nil
527
528
}

529
530
531
532
533
534
func fromCompleteRequest(r CompletionRequest) (api.GenerateRequest, error) {
	options := make(map[string]any)

	switch stop := r.Stop.(type) {
	case string:
		options["stop"] = []string{stop}
535
536
537
538
539
540
541
542
	case []any:
		var stops []string
		for _, s := range stop {
			if str, ok := s.(string); ok {
				stops = append(stops, str)
			} else {
				return api.GenerateRequest{}, fmt.Errorf("invalid type for 'stop' field: %T", s)
			}
543
		}
544
		options["stop"] = stops
545
546
547
548
549
550
551
	}

	if r.MaxTokens != nil {
		options["num_predict"] = *r.MaxTokens
	}

	if r.Temperature != nil {
552
		options["temperature"] = *r.Temperature
553
554
555
556
557
558
559
560
	} else {
		options["temperature"] = 1.0
	}

	if r.Seed != nil {
		options["seed"] = *r.Seed
	}

561
	options["frequency_penalty"] = r.FrequencyPenalty
562

563
	options["presence_penalty"] = r.PresencePenalty
564
565
566
567
568
569
570
571
572
573
574
575

	if r.TopP != 0.0 {
		options["top_p"] = r.TopP
	} else {
		options["top_p"] = 1.0
	}

	return api.GenerateRequest{
		Model:   r.Model,
		Prompt:  r.Prompt,
		Options: options,
		Stream:  &r.Stream,
576
		Suffix:  r.Suffix,
577
578
579
	}, nil
}

580
581
582
583
584
type BaseWriter struct {
	gin.ResponseWriter
}

type ChatWriter struct {
585
586
587
	stream        bool
	streamOptions *StreamOptions
	id            string
588
	BaseWriter
589
590
}

591
type CompleteWriter struct {
592
593
594
	stream        bool
	streamOptions *StreamOptions
	id            string
595
596
597
	BaseWriter
}

598
599
600
601
602
603
604
605
606
type ListWriter struct {
	BaseWriter
}

type RetrieveWriter struct {
	BaseWriter
	model string
}

607
608
609
610
611
type EmbedWriter struct {
	BaseWriter
	model string
}

612
func (w *BaseWriter) writeError(data []byte) (int, error) {
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
	var serr api.StatusError
	err := json.Unmarshal(data, &serr)
	if err != nil {
		return 0, err
	}

	w.ResponseWriter.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w.ResponseWriter).Encode(NewError(http.StatusInternalServerError, serr.Error()))
	if err != nil {
		return 0, err
	}

	return len(data), nil
}

628
func (w *ChatWriter) writeResponse(data []byte) (int, error) {
629
630
631
632
633
634
635
636
	var chatResponse api.ChatResponse
	err := json.Unmarshal(data, &chatResponse)
	if err != nil {
		return 0, err
	}

	// chat chunk
	if w.stream {
637
638
		c := toChunk(w.id, chatResponse)
		d, err := json.Marshal(c)
639
640
641
642
643
644
645
646
647
648
649
		if err != nil {
			return 0, err
		}

		w.ResponseWriter.Header().Set("Content-Type", "text/event-stream")
		_, err = w.ResponseWriter.Write([]byte(fmt.Sprintf("data: %s\n\n", d)))
		if err != nil {
			return 0, err
		}

		if chatResponse.Done {
650
651
652
653
654
655
656
657
658
659
660
661
662
			if w.streamOptions != nil && w.streamOptions.IncludeUsage {
				u := toUsage(chatResponse)
				c.Usage = &u
				c.Choices = []ChunkChoice{}
				d, err := json.Marshal(c)
				if err != nil {
					return 0, err
				}
				_, err = w.ResponseWriter.Write([]byte(fmt.Sprintf("data: %s\n\n", d)))
				if err != nil {
					return 0, err
				}
			}
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
			_, err = w.ResponseWriter.Write([]byte("data: [DONE]\n\n"))
			if err != nil {
				return 0, err
			}
		}

		return len(data), nil
	}

	// chat completion
	w.ResponseWriter.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w.ResponseWriter).Encode(toChatCompletion(w.id, chatResponse))
	if err != nil {
		return 0, err
	}

	return len(data), nil
}

682
func (w *ChatWriter) Write(data []byte) (int, error) {
683
684
	code := w.ResponseWriter.Status()
	if code != http.StatusOK {
685
		return w.writeError(data)
686
687
688
689
690
	}

	return w.writeResponse(data)
}

691
692
693
694
695
696
697
698
699
func (w *CompleteWriter) writeResponse(data []byte) (int, error) {
	var generateResponse api.GenerateResponse
	err := json.Unmarshal(data, &generateResponse)
	if err != nil {
		return 0, err
	}

	// completion chunk
	if w.stream {
700
701
702
703
704
		c := toCompleteChunk(w.id, generateResponse)
		if w.streamOptions != nil && w.streamOptions.IncludeUsage {
			c.Usage = &Usage{}
		}
		d, err := json.Marshal(c)
705
706
707
708
709
710
711
712
713
714
715
		if err != nil {
			return 0, err
		}

		w.ResponseWriter.Header().Set("Content-Type", "text/event-stream")
		_, err = w.ResponseWriter.Write([]byte(fmt.Sprintf("data: %s\n\n", d)))
		if err != nil {
			return 0, err
		}

		if generateResponse.Done {
716
717
718
719
720
721
722
723
724
725
726
727
728
			if w.streamOptions != nil && w.streamOptions.IncludeUsage {
				u := toUsageGenerate(generateResponse)
				c.Usage = &u
				c.Choices = []CompleteChunkChoice{}
				d, err := json.Marshal(c)
				if err != nil {
					return 0, err
				}
				_, err = w.ResponseWriter.Write([]byte(fmt.Sprintf("data: %s\n\n", d)))
				if err != nil {
					return 0, err
				}
			}
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
			_, err = w.ResponseWriter.Write([]byte("data: [DONE]\n\n"))
			if err != nil {
				return 0, err
			}
		}

		return len(data), nil
	}

	// completion
	w.ResponseWriter.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w.ResponseWriter).Encode(toCompletion(w.id, generateResponse))
	if err != nil {
		return 0, err
	}

	return len(data), nil
}

func (w *CompleteWriter) Write(data []byte) (int, error) {
	code := w.ResponseWriter.Status()
	if code != http.StatusOK {
751
		return w.writeError(data)
752
753
754
755
756
	}

	return w.writeResponse(data)
}

757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
func (w *ListWriter) writeResponse(data []byte) (int, error) {
	var listResponse api.ListResponse
	err := json.Unmarshal(data, &listResponse)
	if err != nil {
		return 0, err
	}

	w.ResponseWriter.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w.ResponseWriter).Encode(toListCompletion(listResponse))
	if err != nil {
		return 0, err
	}

	return len(data), nil
}

func (w *ListWriter) Write(data []byte) (int, error) {
	code := w.ResponseWriter.Status()
	if code != http.StatusOK {
776
		return w.writeError(data)
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
	}

	return w.writeResponse(data)
}

func (w *RetrieveWriter) writeResponse(data []byte) (int, error) {
	var showResponse api.ShowResponse
	err := json.Unmarshal(data, &showResponse)
	if err != nil {
		return 0, err
	}

	// retrieve completion
	w.ResponseWriter.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w.ResponseWriter).Encode(toModel(showResponse, w.model))
	if err != nil {
		return 0, err
	}

	return len(data), nil
}

func (w *RetrieveWriter) Write(data []byte) (int, error) {
	code := w.ResponseWriter.Status()
	if code != http.StatusOK {
802
		return w.writeError(data)
803
804
805
806
807
	}

	return w.writeResponse(data)
}

808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
func (w *EmbedWriter) writeResponse(data []byte) (int, error) {
	var embedResponse api.EmbedResponse
	err := json.Unmarshal(data, &embedResponse)
	if err != nil {
		return 0, err
	}

	w.ResponseWriter.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w.ResponseWriter).Encode(toEmbeddingList(w.model, embedResponse))
	if err != nil {
		return 0, err
	}

	return len(data), nil
}

func (w *EmbedWriter) Write(data []byte) (int, error) {
	code := w.ResponseWriter.Status()
	if code != http.StatusOK {
827
		return w.writeError(data)
828
829
830
831
832
	}

	return w.writeResponse(data)
}

833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
func ListMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		w := &ListWriter{
			BaseWriter: BaseWriter{ResponseWriter: c.Writer},
		}

		c.Writer = w

		c.Next()
	}
}

func RetrieveMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(api.ShowRequest{Name: c.Param("model")}); err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, NewError(http.StatusInternalServerError, err.Error()))
			return
		}

		c.Request.Body = io.NopCloser(&b)

		// response writer
		w := &RetrieveWriter{
			BaseWriter: BaseWriter{ResponseWriter: c.Writer},
			model:      c.Param("model"),
		}

		c.Writer = w

		c.Next()
	}
}

867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
func CompletionsMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		var req CompletionRequest
		err := c.ShouldBindJSON(&req)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusBadRequest, NewError(http.StatusBadRequest, err.Error()))
			return
		}

		var b bytes.Buffer
		genReq, err := fromCompleteRequest(req)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusBadRequest, NewError(http.StatusBadRequest, err.Error()))
			return
		}

		if err := json.NewEncoder(&b).Encode(genReq); err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, NewError(http.StatusInternalServerError, err.Error()))
			return
		}

		c.Request.Body = io.NopCloser(&b)

		w := &CompleteWriter{
891
892
893
894
			BaseWriter:    BaseWriter{ResponseWriter: c.Writer},
			stream:        req.Stream,
			id:            fmt.Sprintf("cmpl-%d", rand.Intn(999)),
			streamOptions: req.StreamOptions,
895
896
		}

897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
		c.Writer = w
		c.Next()
	}
}

func EmbeddingsMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		var req EmbedRequest
		err := c.ShouldBindJSON(&req)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusBadRequest, NewError(http.StatusBadRequest, err.Error()))
			return
		}

		if req.Input == "" {
			req.Input = []string{""}
		}

		if req.Input == nil {
			c.AbortWithStatusJSON(http.StatusBadRequest, NewError(http.StatusBadRequest, "invalid input"))
			return
		}

		if v, ok := req.Input.([]any); ok && len(v) == 0 {
			c.AbortWithStatusJSON(http.StatusBadRequest, NewError(http.StatusBadRequest, "invalid input"))
			return
		}

		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(api.EmbedRequest{Model: req.Model, Input: req.Input}); err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, NewError(http.StatusInternalServerError, err.Error()))
			return
		}

		c.Request.Body = io.NopCloser(&b)

		w := &EmbedWriter{
			BaseWriter: BaseWriter{ResponseWriter: c.Writer},
			model:      req.Model,
		}

938
939
940
941
942
943
		c.Writer = w

		c.Next()
	}
}

944
func ChatMiddleware() gin.HandlerFunc {
945
946
947
948
949
950
951
952
953
954
955
956
957
958
	return func(c *gin.Context) {
		var req ChatCompletionRequest
		err := c.ShouldBindJSON(&req)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusBadRequest, NewError(http.StatusBadRequest, err.Error()))
			return
		}

		if len(req.Messages) == 0 {
			c.AbortWithStatusJSON(http.StatusBadRequest, NewError(http.StatusBadRequest, "[] is too short - 'messages'"))
			return
		}

		var b bytes.Buffer
959
960
961
962

		chatReq, err := fromChatRequest(req)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusBadRequest, NewError(http.StatusBadRequest, err.Error()))
963
			return
964
965
966
		}

		if err := json.NewEncoder(&b).Encode(chatReq); err != nil {
967
968
969
970
971
972
			c.AbortWithStatusJSON(http.StatusInternalServerError, NewError(http.StatusInternalServerError, err.Error()))
			return
		}

		c.Request.Body = io.NopCloser(&b)

973
		w := &ChatWriter{
974
975
976
977
			BaseWriter:    BaseWriter{ResponseWriter: c.Writer},
			stream:        req.Stream,
			id:            fmt.Sprintf("chatcmpl-%d", rand.Intn(999)),
			streamOptions: req.StreamOptions,
978
979
980
981
982
983
984
		}

		c.Writer = w

		c.Next()
	}
}