add cf

.gitignore

@@ -33,4 +33,5 @@ bin/
 # Develop tools
 .vscode/
 .idea/
-*.swp
+*.swp
+config.json

README.md

@@ -58,6 +58,7 @@
 3. 我目前就试了下 `VSCode` ，至于 `JetBrains` 等IDE尚未适配，如果你有相关经验，请告诉我。
 4. 项目基于 `MIT` 协议发布，你可以修改，请保留原作者信息。
 5. 有什么问题，请在论坛 https://linux.do 讨论，欢迎PR。
+6. 可以支持cf，目前最大token 1500左右 配置里字段名是 `"codex_model_default":"@hf/thebloke/deepseek-coder-6.7b-instruct-awq"`
 
 ### Star History
 

config.json

@@ -2,11 +2,12 @@
   "bind": "127.0.0.1:8181",
   "proxy_url": "",
   "timeout": 600,
-  "codex_api_base": "https://api-proxy.oaipro.com/v1/completions",
-  "codex_api_key": "sk-xxx",
+  "codex_api_base": "https://api.oaipro.com/v1/completions",
+  "codex_api_key": "sk-4VTExcjhhlexSBWXCfDd15A102094580B9616062E0Cb6319",
   "codex_api_organization": "",
   "codex_api_project": "",
   "codex_max_tokens": 4093,
+  "codex_model_default":"gpt-3.5-turbo-instruct",
   "chat_api_base": "https://api-proxy.oaipro.com/v1",
   "chat_api_key": "sk-xxx",
   "chat_api_organization": "",

main.go

@@ -1,10 +1,12 @@
 package main
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"github.com/gin-gonic/gin"
 	"github.com/linux-do/tiktoken-go"
 	"github.com/tidwall/gjson"
@@ -23,6 +25,129 @@ import (
 
 const INSTRUCT_MODEL = "gpt-3.5-turbo-instruct"
 
+type GPTMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+type StreamResponse struct {
+	Response string `json:"response"`
+}
+type Message struct {
+	Role    string  `json:"role,omitempty"`
+	Content any     `json:"content,omitempty"`
+	Name    *string `json:"name,omitempty"`
+}
+type ChatCompletionsStreamResponseChoice struct {
+	Index        int     `json:"index"`
+	Delta        Message `json:"delta"`
+	FinishReason *string `json:"finish_reason,omitempty"`
+}
+
+type ChatCompletionsStreamResponse struct {
+	Id      string                                `json:"id"`
+	Object  string                                `json:"object"`
+	Created int64                                 `json:"created"`
+	Model   string                                `json:"model"`
+	Choices []ChatCompletionsStreamResponseChoice `json:"choices"`
+}
+type CustomEvent struct {
+	Event string
+	Id    string
+	Retry uint
+	Data  interface{}
+}
+type stringWriter interface {
+	io.Writer
+	writeString(string) (int, error)
+}
+
+type stringWrapper struct {
+	io.Writer
+}
+
+var dataReplacer = strings.NewReplacer(
+	"\n", "\ndata:",
+	"\r", "\\r")
+var contentType = []string{"text/event-stream"}
+var noCache = []string{"no-cache"}
+
+func (w stringWrapper) writeString(str string) (int, error) {
+	return w.Writer.Write([]byte(str))
+}
+func checkWriter(writer io.Writer) stringWriter {
+	if w, ok := writer.(stringWriter); ok {
+		return w
+	} else {
+		return stringWrapper{writer}
+	}
+}
+func encode(writer io.Writer, event CustomEvent) error {
+	w := checkWriter(writer)
+	return writeData(w, event.Data)
+}
+func writeData(w stringWriter, data interface{}) error {
+	dataReplacer.WriteString(w, fmt.Sprint(data))
+	if strings.HasPrefix(data.(string), "data") {
+		w.writeString("\n\n")
+	}
+	return nil
+}
+func (r CustomEvent) Render(w http.ResponseWriter) error {
+	r.WriteContentType(w)
+	return encode(w, r)
+}
+
+func (r CustomEvent) WriteContentType(w http.ResponseWriter) {
+	header := w.Header()
+	header["Content-Type"] = contentType
+
+	if _, exist := header["Cache-Control"]; !exist {
+		header["Cache-Control"] = noCache
+	}
+}
+
+func GetTimestamp() int64 {
+	return time.Now().Unix()
+}
+func StreamResponseCloudflare2OpenAI(cloudflareResponse *StreamResponse) *ChatCompletionsStreamResponse {
+	var choice ChatCompletionsStreamResponseChoice
+	choice.Delta.Content = cloudflareResponse.Response
+	choice.Delta.Role = "assistant"
+	openaiResponse := ChatCompletionsStreamResponse{
+		Object:  "chat.completion.chunk",
+		Choices: []ChatCompletionsStreamResponseChoice{choice},
+		Created: GetTimestamp(),
+	}
+	return &openaiResponse
+}
+func StreamResponse2OpenAI(cloudflareResponse *StreamResponse) *ChatCompletionsStreamResponse {
+	var choice ChatCompletionsStreamResponseChoice
+	choice.Delta.Content = cloudflareResponse.Response
+	choice.Delta.Role = "assistant"
+	openaiResponse := ChatCompletionsStreamResponse{
+		Object:  "chat.completion.chunk",
+		Choices: []ChatCompletionsStreamResponseChoice{choice},
+		Created: GetTimestamp(),
+	}
+	return &openaiResponse
+}
+func SetEventStreamHeaders(c *gin.Context) {
+	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	c.Writer.Header().Set("Cache-Control", "no-cache")
+	c.Writer.Header().Set("Connection", "keep-alive")
+	c.Writer.Header().Set("Transfer-Encoding", "chunked")
+	c.Writer.Header().Set("X-Accel-Buffering", "no")
+}
+
+const (
+	RequestIdKey = "X-Oneapi-Request-Id"
+)
+
+func GetResponseID(c *gin.Context) string {
+	logID := c.GetString(RequestIdKey)
+	return fmt.Sprintf("chatcmpl-%s", logID)
+}
+
 type config struct {
 	Bind                 string            `json:"bind"`
 	ProxyUrl             string            `json:"proxy_url"`
@@ -32,6 +157,7 @@ type config struct {
 	CodexApiOrganization string            `json:"codex_api_organization"`
 	CodexApiProject      string            `json:"codex_api_project"`
 	CodexMaxTokens       int               `json:"codex_max_tokens"`
+	CodexModelDefault    string            `json:"codex_model_default"`
 	ChatApiBase          string            `json:"chat_api_base"`
 	ChatApiKey           string            `json:"chat_api_key"`
 	ChatApiOrganization  string            `json:"chat_api_organization"`
@@ -268,9 +394,33 @@ func (s *ProxyService) codeCompletions(c *gin.Context) {
 
 	body, _ = sjson.DeleteBytes(body, "extra")
 	body, _ = sjson.DeleteBytes(body, "nwo")
-	body, _ = sjson.SetBytes(body, "model", INSTRUCT_MODEL)
-
-	proxyUrl := s.cfg.CodexApiBase + "/completions"
+	var model string
+	proxyUrl := s.cfg.CodexApiBase
+	if s.cfg.CodexModelDefault == "" || s.cfg.CodexModelDefault == INSTRUCT_MODEL {
+		model = INSTRUCT_MODEL
+		proxyUrl = proxyUrl + "/completions"
+	} else {
+		model = s.cfg.CodexModelDefault
+	}
+	body, _ = sjson.SetBytes(body, "model", model)
+	if model == "deepseek-coder" {
+		message := gjson.GetBytes(body, "prompt").String()
+		body, _ = sjson.DeleteBytes(body, "prompt")
+		msg := make([]GPTMessage, 0)
+		msg = append(msg, GPTMessage{Role: "system", Content: "You are a helpful assistant"})
+		msg = append(msg, GPTMessage{Role: "user", Content: message})
+		body, _ = sjson.SetBytes(body, "messages", msg)
+		body, _ = sjson.DeleteBytes(body, "n")
+	} else if strings.HasPrefix(model, "@") {
+		proxyUrl = s.cfg.CodexApiBase
+		message := gjson.GetBytes(body, "prompt").String()
+		body, _ = sjson.DeleteBytes(body, "prompt")
+		msg := make([]GPTMessage, 0)
+		msg = append(msg, GPTMessage{Role: "system", Content: ""})
+		msg = append(msg, GPTMessage{Role: "user", Content: message})
+		body, _ = sjson.SetBytes(body, "messages", msg)
+		body, _ = sjson.DeleteBytes(body, "n")
+	}
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, proxyUrl, io.NopCloser(bytes.NewBuffer(body)))
 	if nil != err {
 		abortCodex(c, http.StatusInternalServerError)
@@ -314,7 +464,74 @@ func (s *ProxyService) codeCompletions(c *gin.Context) {
 		c.Header("Content-Type", contentType)
 	}
 
-	_, _ = io.Copy(c.Writer, resp.Body)
+	scanner := bufio.NewScanner(resp.Body)
+	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
+		if atEOF && len(data) == 0 {
+			return 0, nil, nil
+		}
+		if i := bytes.IndexByte(data, '\n'); i >= 0 {
+			return i + 1, data[0:i], nil
+		}
+		if atEOF {
+			return len(data), data, nil
+		}
+		return 0, nil, nil
+	})
+	dataChan := make(chan string)
+	stopChan := make(chan bool)
+	go func() {
+		for scanner.Scan() {
+			data := scanner.Text()
+			if len(data) < len("data: ") {
+				continue
+			}
+			data = strings.TrimPrefix(data, "data: ")
+			dataChan <- data
+		}
+		stopChan <- true
+	}()
+	SetEventStreamHeaders(c)
+	id := GetResponseID(c)
+	responseModel := c.GetString("original_model")
+	var responseText string
+	c.Stream(func(w io.Writer) bool {
+		select {
+		case data := <-dataChan:
+			// some implementations may add \r at the end of data
+			data = strings.TrimSuffix(data, "\r")
+			var codeResponse StreamResponse
+			err := json.Unmarshal([]byte(data), &codeResponse)
+			if err != nil {
+				if data == "[DONE]" {
+					return true
+				}
+				log.Println("error unmarshalling stream response: ", err.Error())
+				return true
+			}
+			if model != INSTRUCT_MODEL {
+				response := StreamResponseCloudflare2OpenAI(&codeResponse)
+				if response == nil {
+					return true
+				}
+				responseText += codeResponse.Response
+				response.Id = id
+				response.Model = responseModel
+				jsonStr, err := json.Marshal(response)
+				if err != nil {
+					log.Println("error marshalling stream response: ", err.Error())
+					return true
+				}
+				c.Render(-1, CustomEvent{Data: "data: " + string(jsonStr)})
+			} else {
+				c.Render(-1, CustomEvent{Data: "data:" + string(data)})
+			}
+			return true
+		case <-stopChan:
+			c.Render(-1, CustomEvent{Data: "data: [DONE]"})
+			return false
+		}
+	})
+	_ = resp.Body.Close()
 }
 
 func main() {