9um3yhdu commited on
Commit
e64e819
·
verified ·
1 Parent(s): 0ef4d41

Upload 2 files

Browse files
Files changed (2) hide show
  1. Dockerfile +8 -0
  2. main.go +311 -0
Dockerfile ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ FROM golang:1.23-bullseye
2
+ WORKDIR /app
3
+ RUN mkdir -p /.cache && \
4
+ chmod -R 777 /.cache
5
+ COPY main.go ./
6
+ RUN uname -a
7
+ EXPOSE 7860
8
+ CMD ["go", "run", "main.go"]
main.go ADDED
@@ -0,0 +1,311 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ package main
2
+
3
+ import (
4
+ "bufio"
5
+ "bytes"
6
+ "encoding/json"
7
+ "fmt"
8
+ "io/ioutil"
9
+ "log"
10
+ "net/http"
11
+ "strings"
12
+ "sync"
13
+ "time"
14
+ )
15
+
16
+ type OpenAIRequest struct {
17
+ Model string `json:"model"`
18
+ Messages []Message `json:"messages"`
19
+ Stream bool `json:"stream"`
20
+ }
21
+
22
+ type Message struct {
23
+ Role string `json:"role"`
24
+ Content string `json:"content"`
25
+ }
26
+
27
+ type DeepSeekResponse struct {
28
+ Code int `json:"code"`
29
+ Msg string `json:"msg"`
30
+ Message string `json:"message"`
31
+ APISource string `json:"api_source"`
32
+ }
33
+
34
+ type OpenAIResponse struct {
35
+ ID string `json:"id"`
36
+ Object string `json:"object"`
37
+ Created int64 `json:"created"`
38
+ Model string `json:"model"`
39
+ Choices []Choice `json:"choices"`
40
+ }
41
+
42
+ type Choice struct {
43
+ Index int `json:"index"`
44
+ Message Message `json:"message"`
45
+ FinishReason string `json:"finish_reason"`
46
+ }
47
+
48
+ type StreamChoice struct {
49
+ Delta StreamMessage `json:"delta"`
50
+ Index int `json:"index"`
51
+ FinishReason *string `json:"finish_reason"`
52
+ }
53
+
54
+ type StreamMessage struct {
55
+ Role string `json:"role,omitempty"`
56
+ Content string `json:"content,omitempty"`
57
+ }
58
+
59
+ type StreamResponse struct {
60
+ ID string `json:"id"`
61
+ Object string `json:"object"`
62
+ Created int64 `json:"created"`
63
+ Model string `json:"model"`
64
+ Choices []StreamChoice `json:"choices"`
65
+ }
66
+
67
+ var (
68
+ requestCount int64
69
+ requestLog []string
70
+ lastMinute time.Time
71
+ rpm int
72
+ logMutex sync.Mutex
73
+ )
74
+
75
+ func init() {
76
+ lastMinute = time.Now()
77
+ requestLog = make([]string, 0, 10000)
78
+ }
79
+
80
+ func main() {
81
+ http.HandleFunc("/", handleStats)
82
+ http.HandleFunc("/log", handleLogs)
83
+ http.HandleFunc("/hf/v1/chat/completions", handleChat)
84
+ log.Fatal(http.ListenAndServe(":7860", nil))
85
+ }
86
+
87
+ func handleStats(w http.ResponseWriter, r *http.Request) {
88
+ if r.URL.Path != "/" {
89
+ http.NotFound(w, r)
90
+ return
91
+ }
92
+
93
+ logMutex.Lock()
94
+ currentRPM := rpm
95
+ totalRequests := requestCount
96
+ logMutex.Unlock()
97
+
98
+ fmt.Fprintf(w, "总请求次数: %d\n每分钟请求数(RPM): %d", totalRequests, currentRPM)
99
+ }
100
+
101
+ func handleLogs(w http.ResponseWriter, r *http.Request) {
102
+ auth := r.URL.Query().Get("auth")
103
+ if auth != "smnet" {
104
+ http.Error(w, "未授权访问", http.StatusUnauthorized)
105
+ return
106
+ }
107
+
108
+ logMutex.Lock()
109
+ logs := make([]string, len(requestLog))
110
+ copy(logs, requestLog)
111
+ logMutex.Unlock()
112
+
113
+ for _, log := range logs {
114
+ fmt.Fprintln(w, log)
115
+ }
116
+ }
117
+
118
+ func handleChat(w http.ResponseWriter, r *http.Request) {
119
+ logMutex.Lock()
120
+ requestCount++
121
+ now := time.Now()
122
+ if now.Sub(lastMinute) >= time.Minute {
123
+ rpm = 1
124
+ lastMinute = now
125
+ } else {
126
+ rpm++
127
+ }
128
+
129
+ clientIP := r.Header.Get("X-Real-IP")
130
+ if clientIP == "" {
131
+ clientIP = r.Header.Get("X-Forwarded-For")
132
+ if clientIP == "" {
133
+ clientIP = r.RemoteAddr
134
+ }
135
+ }
136
+
137
+ logEntry := fmt.Sprintf("[%s] IP:%s 新请求处理", now.Format("2006-01-02 15:04:05"), clientIP)
138
+ if len(requestLog) >= 5000 {
139
+ requestLog = requestLog[1:]
140
+ }
141
+ requestLog = append(requestLog, logEntry)
142
+ logMutex.Unlock()
143
+
144
+ if r.Method != http.MethodPost {
145
+ log.Printf("错误: 不支持的请求方法 %s", r.Method)
146
+ http.Error(w, "仅支持 POST 请求", http.StatusMethodNotAllowed)
147
+ return
148
+ }
149
+
150
+ body, err := ioutil.ReadAll(r.Body)
151
+ if err != nil {
152
+ log.Printf("错误: 读取请求失败 - %v", err)
153
+ http.Error(w, "读取请求失败", http.StatusBadRequest)
154
+ return
155
+ }
156
+
157
+ var openAIReq OpenAIRequest
158
+ if err := json.Unmarshal(body, &openAIReq); err != nil {
159
+ log.Printf("错误: 请求格式错误 - %v", err)
160
+ http.Error(w, "请求格式错误", http.StatusBadRequest)
161
+ return
162
+ }
163
+
164
+ log.Printf("用户问题: %s", openAIReq.Messages[len(openAIReq.Messages)-1].Content)
165
+
166
+ var apiURL string
167
+ var modelName string
168
+ switch openAIReq.Model {
169
+ case "deepseek-r1":
170
+ apiURL = "https://api.deepinfra.com/v1/openai/chat/completions"
171
+ modelName = "deepseek-ai/DeepSeek-R1"
172
+ default:
173
+ apiURL = "https://api.deepinfra.com/v1/openai/chat/completions"
174
+ modelName = "deepseek-ai/DeepSeek-V3"
175
+ }
176
+
177
+ deepseekReq := map[string]interface{}{
178
+ "messages": openAIReq.Messages,
179
+ "stream": true,
180
+ "model": modelName,
181
+ }
182
+
183
+ deepseekBody, err := json.Marshal(deepseekReq)
184
+ if err != nil {
185
+ log.Printf("错误: 构造请求失败 - %v", err)
186
+ http.Error(w, "构造请求失败", http.StatusInternalServerError)
187
+ return
188
+ }
189
+
190
+ maxRetries := 200
191
+ var tryRequest func() (string, error)
192
+
193
+ tryRequest = func() (string, error) {
194
+ req, err := http.NewRequest("POST", apiURL, bytes.NewBuffer(deepseekBody))
195
+ if err != nil {
196
+ return "", fmt.Errorf("创建请求失败: %v", err)
197
+ }
198
+ req.Header.Set("Content-Type", "application/json")
199
+
200
+ client := &http.Client{}
201
+ resp, err := client.Do(req)
202
+ if err != nil {
203
+ return "", fmt.Errorf("请求失败: %v", err)
204
+ }
205
+ defer resp.Body.Close()
206
+
207
+ var fullMessage string
208
+ scanner := bufio.NewScanner(resp.Body)
209
+
210
+ for scanner.Scan() {
211
+ line := scanner.Text()
212
+ if openAIReq.Stream {
213
+ _, err = fmt.Fprintf(w, "%s\n", line)
214
+ if err != nil {
215
+ return "", fmt.Errorf("写入流式响应失败: %v", err)
216
+ }
217
+ w.(http.Flusher).Flush()
218
+ }
219
+
220
+ if !strings.HasPrefix(line, "data: ") {
221
+ continue
222
+ }
223
+
224
+ data := strings.TrimPrefix(line, "data: ")
225
+ if data == "[DONE]" {
226
+ break
227
+ }
228
+
229
+ var streamResp StreamResponse
230
+ if err := json.Unmarshal([]byte(data), &streamResp); err != nil {
231
+ continue
232
+ }
233
+
234
+ if len(streamResp.Choices) > 0 && streamResp.Choices[0].Delta.Content != "" {
235
+ fullMessage += streamResp.Choices[0].Delta.Content
236
+ }
237
+ }
238
+
239
+ if fullMessage == "" {
240
+ return "", fmt.Errorf("收到空回复")
241
+ }
242
+
243
+ return fullMessage, nil
244
+ }
245
+
246
+ var fullMessage string
247
+ var lastError error
248
+
249
+ for i := 0; i < maxRetries; i++ {
250
+ fullMessage, lastError = tryRequest()
251
+ if lastError == nil {
252
+ break
253
+ }
254
+ log.Printf("第 %d 次尝试失败: %v,准备重试", i+1, lastError)
255
+ time.Sleep(time.Second * time.Duration(i+1))
256
+ }
257
+
258
+ if lastError != nil {
259
+ log.Printf("错误: 所有重试都失败 - %v", lastError)
260
+ http.Error(w, "服务暂时不可用", http.StatusServiceUnavailable)
261
+ return
262
+ }
263
+
264
+ log.Printf("AI回答: %s", fullMessage)
265
+
266
+ if !openAIReq.Stream {
267
+ openAIResp := OpenAIResponse{
268
+ ID: "chatcmpl-" + generateRandomString(10),
269
+ Object: "chat.completion",
270
+ Created: getCurrentTimestamp(),
271
+ Model: openAIReq.Model,
272
+ Choices: []Choice{
273
+ {
274
+ Index: 0,
275
+ Message: Message{
276
+ Role: "assistant",
277
+ Content: fullMessage,
278
+ },
279
+ FinishReason: "stop",
280
+ },
281
+ },
282
+ }
283
+
284
+ w.Header().Set("Content-Type", "application/json")
285
+ json.NewEncoder(w).Encode(openAIResp)
286
+ }
287
+ }
288
+
289
+ func generateRandomString(length int) string {
290
+ return "SomeApiResponse"
291
+ }
292
+
293
+ func getCurrentTimestamp() int64 {
294
+ return time.Now().Unix()
295
+ }
296
+
297
+ func writeSSE(w http.ResponseWriter, data interface{}) error {
298
+ jsonData, err := json.Marshal(data)
299
+ if err != nil {
300
+ http.Error(w, "JSON编码失败", http.StatusInternalServerError)
301
+ return err
302
+ }
303
+
304
+ _, err = fmt.Fprintf(w, "data: %s\n\n", jsonData)
305
+ if err != nil {
306
+ http.Error(w, "写入响应失败", http.StatusInternalServerError)
307
+ return err
308
+ }
309
+
310
+ return nil
311
+ }