From 2092b17b2f78dba392e319b3930a7abf57786e67 Mon Sep 17 00:00:00 2001
From: zhuiyue132 <fenda.app@qq.com>
Date: Fri, 11 Oct 2024 16:50:58 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=20siliconflow=20fim?=
 =?UTF-8?q?=20=E7=9A=84=E6=94=AF=E6=8C=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

siliconflow 已支持标准格式的 FIM 补全,特此 PR
---
 README.md | 15 +++++++++++++++
 main.go   | 21 +++++++++++++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index dd0ce6e..734b68b 100644
--- a/README.md
+++ b/README.md
@@ -69,6 +69,21 @@
   "code_instruct_model": "deepseek-coder",
 ```
 
+### Siliconflow 设置
+如果你希望使用 Siliconflow FIM 模型来进行代码补全，着重修改以下配置：
+
+```json
+  "codex_api_base": "https://api.siliconflow.cn/v1",
+  "codex_api_key": "sk-xxx,sk-xxx2,sk-xxx3...",
+  "code_instruct_model": "Qwen/Qwen2.5-Coder-7B-Instruct",
+```
+
+截至目前，Siliconflow 共有三个模型支持 FIM。分别是 `Qwen/Qwen2.5-Coder-7B-Instruct`、`deepseek-ai/DeepSeek-Coder-V2-Instruct` 、`deepseek-ai/DeepSeek-V2.5`。其中 `Qwen/Qwen2.5-Coder-7B-Instruct` 是免费模型，另外两个是收费模型。
+
+如果你有很多 Siliconflow API Key, 可以以英文逗号分隔填入`codex_api_key`字段, 这样可以很好的避免Siliconflow官方的 TPM RateLimit 对你编码速度影响(尤其使用收费模型时，用户级别较低，TPM 最低只有 10k)。
+
+
+
 ### 本地大模型设置
 1. 安装ollama 
 2. ollama run stable-code:code  (这个模型较小，大部分显卡都能跑)  
diff --git a/main.go b/main.go
index 78f8215..97d60e0 100644
--- a/main.go
+++ b/main.go
@@ -19,6 +19,7 @@ import (
 	"strconv"
 	"strings"
 	"time"
+    "math/rand"
 )
 
 const DefaultInstructModel = "gpt-3.5-turbo-instruct"
@@ -27,6 +28,8 @@ const StableCodeModelPrefix = "stable-code"
 
 const DeepSeekCoderModel = "deepseek-coder"
 
+var SiliconflowModels = []string{"deepseek-ai/DeepSeek-V2.5", "deepseek-ai/DeepSeek-Coder-V2-Instruct", "Qwen/Qwen2.5-Coder-7B-Instruct"}
+
 type config struct {
 	Bind                 string            `json:"bind"`
 	ProxyUrl             string            `json:"proxy_url"`
@@ -474,6 +477,10 @@ func (s *ProxyService) completions(c *gin.Context) {
 	_, _ = io.Copy(c.Writer, resp.Body)
 }
 
+func contains(arr []string, str string) bool {
+    return strings.Contains(strings.Join(arr, ","), str)
+}
+
 func (s *ProxyService) codeCompletions(c *gin.Context) {
 	ctx := c.Request.Context()
 
@@ -499,7 +506,7 @@ func (s *ProxyService) codeCompletions(c *gin.Context) {
 	}
 
 	req.Header.Set("Content-Type", "application/json")
-	req.Header.Set("Authorization", "Bearer "+s.cfg.CodexApiKey)
+	req.Header.Set("Authorization", "Bearer " + getRandomApiKey(s.cfg.CodexApiKey))
 	if "" != s.cfg.CodexApiOrganization {
 		req.Header.Set("OpenAI-Organization", s.cfg.CodexApiOrganization)
 	}
@@ -538,6 +545,16 @@ func (s *ProxyService) codeCompletions(c *gin.Context) {
 	_, _ = io.Copy(c.Writer, resp.Body)
 }
 
+// 随机取一个apiKey
+func getRandomApiKey(paramStr string) string {
+    params := strings.Split(paramStr, ",")
+    rand.Seed(time.Now().UnixNano())
+    randomIndex := rand.Intn(len(params))
+	fmt.Println("Code completion API Key index:", randomIndex)
+	fmt.Println("Code completion API Key:", strings.TrimSpace(params[randomIndex]))
+    return strings.TrimSpace(params[randomIndex])
+}
+
 func ConstructRequestBody(body []byte, cfg *config) []byte {
 	body, _ = sjson.DeleteBytes(body, "extra")
 	body, _ = sjson.DeleteBytes(body, "nwo")
@@ -549,7 +566,7 @@ func ConstructRequestBody(body []byte, cfg *config) []byte {
 
 	if strings.Contains(cfg.CodeInstructModel, StableCodeModelPrefix) {
 		return constructWithStableCodeModel(body)
-	} else if strings.HasPrefix(cfg.CodeInstructModel, DeepSeekCoderModel) {
+	} else if strings.HasPrefix(cfg.CodeInstructModel, DeepSeekCoderModel) || contains(SiliconflowModels, cfg.CodeInstructModel) {
 		if gjson.GetBytes(body, "n").Int() > 1 {
 			body, _ = sjson.SetBytes(body, "n", 1)
 		}