Update api/main.go

api/main.go

@@ -1,62 +1,73 @@
 package main
 
 import (
-	"bufio"
-	"encoding/json"
-	"fmt"
-	"net/http"
-	"strings"
-	"time"
+    "bufio"
+    "encoding/json"
+    "fmt"
+    "net/http"
+    "strings"
+    "time"
 )
 
 type YouChatResponse struct {
-	YouChatToken string `json:"youChatToken"`
+    YouChatToken string `json:"youChatToken"`
 }
 
 type OpenAIStreamResponse struct {
-	ID      string   `json:"id"`
-	Object  string   `json:"object"`
-	Created int64    `json:"created"`
-	Model   string   `json:"model"`
-	Choices []Choice `json:"choices"`
+    ID      string   `json:"id"`
+    Object  string   `json:"object"`
+    Created int64    `json:"created"`
+    Model   string   `json:"model"`
+    Choices []Choice `json:"choices"`
 }
 
 type Choice struct {
-	Delta        Delta  `json:"delta"`
-	Index        int    `json:"index"`
-	FinishReason string `json:"finish_reason"`
+    Delta        Delta  `json:"delta"`
+    Index        int    `json:"index"`
+    FinishReason string `json:"finish_reason"`
 }
 
 type Delta struct {
-	Content string `json:"content"`
+    Content string `json:"content"`
 }
 
 type OpenAIRequest struct {
-	Messages []Message `json:"messages"`
-	Stream   bool      `json:"stream"`
-	Model    string    `json:"model"`
+    Messages []Message `json:"messages"`
+    Stream   bool      `json:"stream"`
+    Model    string    `json:"model"`
 }
 
 type Message struct {
-	Role    string `json:"role"`
-	Content string `json:"content"`
+    Role    string `json:"role"`
+    Content string `json:"content"`
 }
 
 type OpenAIResponse struct {
-	ID      string         `json:"id"`
-	Object  string         `json:"object"`
-	Created int64          `json:"created"`
-	Model   string         `json:"model"`
-	Choices []OpenAIChoice `json:"choices"`
+    ID      string         `json:"id"`
+    Object  string         `json:"object"`
+    Created int64          `json:"created"`
+    Model   string         `json:"model"`
+    Choices []OpenAIChoice `json:"choices"`
 }
 
 type OpenAIChoice struct {
-	Message      Message `json:"message"`
-	Index        int     `json:"index"`
-	FinishReason string  `json:"finish_reason"`
+    Message      Message `json:"message"`
+    Index        int     `json:"index"`
+    FinishReason string  `json:"finish_reason"`
+}
+
+type ModelResponse struct {
+    Object string        `json:"object"`
+    Data   []ModelDetail `json:"data"`
+}
+
+type ModelDetail struct {
+    ID      string `json:"id"`
+    Object  string `json:"object"`
+    Created int64  `json:"created"`
+    OwnedBy string `json:"owned_by"`
 }
 
-// modelMap 存储 OpenAI 模型名称到 You.com 模型名称的映射。
 var modelMap = map[string]string{
     "deepseek-reasoner":       "deepseek_r1",
     "deepseek-chat":           "deepseek_v3",
@@ -83,7 +94,7 @@ var modelMap = map[string]string{
     "gemini-2-flash":          "gemini_2_flash",
     "gemini-1.5-flash":        "gemini_1_5_flash",
     "gemini-1.5-pro":          "gemini_1_5_pro",
-    "databricks-dbrx-instruct":"databricks_dbrx_instruct",
+    "databricks-dbrx-instruct": "databricks_dbrx_instruct",
     "qwen-2.5-72b":            "qwen2p5_72b",
     "qwen-2.5-coder-32b":      "qwen2p5_coder_32b",
     "command-r-plus":          "command_r_plus",
@@ -92,256 +103,289 @@ var modelMap = map[string]string{
 }
 
 func getReverseModelMap() map[string]string {
-	reverse := make(map[string]string, len(modelMap))
-	for k, v := range modelMap {
-		reverse[v] = k
-	}
-	return reverse
+    reverse := make(map[string]string, len(modelMap))
+    for k, v := range modelMap {
+        reverse[v] = k
+    }
+    return reverse
 }
 
 func mapModelName(openAIModel string) string {
-	if mappedModel, exists := modelMap[openAIModel]; exists {
-		return mappedModel
-	}
-	return "deepseek_v3"
+    if mappedModel, exists := modelMap[openAIModel]; exists {
+        return mappedModel
+    }
+    return "deepseek_v3"
 }
 
 func reverseMapModelName(youModel string) string {
-	reverseMap := getReverseModelMap()
-	if mappedModel, exists := reverseMap[youModel]; exists {
-		return mappedModel
-	}
-	return "deepseek-chat"
+    reverseMap := getReverseModelMap()
+    if mappedModel, exists := reverseMap[youModel]; exists {
+        return mappedModel
+    }
+    return "deepseek-chat"
 }
 
 var originalModel string
 
 func Handler(w http.ResponseWriter, r *http.Request) {
-	if r.URL.Path != "/hf/v1/chat/completions" {
-		w.Header().Set("Content-Type", "application/json")
-		json.NewEncoder(w).Encode(map[string]string{
-			"status":  "You2Api Service Running...",
-			"message": "MoLoveSze...",
-		})
-		return
-	}
-
-	w.Header().Set("Access-Control-Allow-Origin", "*")
-	w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
-	w.Header().Set("Access-Control-Allow-Headers", "*")
-
-	if r.Method == "OPTIONS" {
-		w.WriteHeader(http.StatusOK)
-		return
-	}
-
-	authHeader := r.Header.Get("Authorization")
-	if !strings.HasPrefix(authHeader, "Bearer ") {
-		http.Error(w, "Missing or invalid authorization header", http.StatusUnauthorized)
-		return
-	}
-	dsToken := strings.TrimPrefix(authHeader, "Bearer ")
-
-	var openAIReq OpenAIRequest
-	if err := json.NewDecoder(r.Body).Decode(&openAIReq); err != nil {
-		http.Error(w, "Invalid request body", http.StatusBadRequest)
-		return
-	}
-
-	originalModel = openAIReq.Model
-	lastMessage := openAIReq.Messages[len(openAIReq.Messages)-1].Content
-	var chatHistory []map[string]interface{}
-	for _, msg := range openAIReq.Messages {
-		chatMsg := map[string]interface{}{
-			"question": msg.Content,
-			"answer":   "",
-		}
-		if msg.Role == "assistant" {
-			chatMsg["question"] = ""
-			chatMsg["answer"] = msg.Content
-		}
-		chatHistory = append(chatHistory, chatMsg)
-	}
-
-	chatHistoryJSON, _ := json.Marshal(chatHistory)
-
-	youReq, _ := http.NewRequest("GET", "https://you.com/api/streamingSearch", nil)
-
-	q := youReq.URL.Query()
-	q.Add("q", lastMessage)
-	q.Add("page", "1")
-	q.Add("count", "10")
-	q.Add("safeSearch", "Moderate")
-	q.Add("mkt", "zh-HK")
-	q.Add("enable_worklow_generation_ux", "true")
-	q.Add("domain", "youchat")
-	q.Add("use_personalization_extraction", "true")
-	q.Add("pastChatLength", fmt.Sprintf("%d", len(chatHistory)-1))
-	q.Add("selectedChatMode", "custom")
-	q.Add("selectedAiModel", mapModelName(openAIReq.Model))
-	q.Add("enable_agent_clarification_questions", "true")
-	q.Add("use_nested_youchat_updates", "true")
-	q.Add("chat", string(chatHistoryJSON))
-	youReq.URL.RawQuery = q.Encode()
-
-	youReq.Header = http.Header{
-		"sec-ch-ua-platform":         {"Windows"},
-		"Cache-Control":              {"no-cache"},
-		"sec-ch-ua":                  {`"Not(A:Brand";v="99", "Microsoft Edge";v="133", "Chromium";v="133"`},
-		"sec-ch-ua-bitness":          {"64"},
-		"sec-ch-ua-model":            {""},
-		"sec-ch-ua-mobile":           {"?0"},
-		"sec-ch-ua-arch":             {"x86"},
-		"sec-ch-ua-full-version":     {"133.0.3065.39"},
-		"Accept":                     {"text/event-stream"},
-		"User-Agent":                 {"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36 Edg/133.0.0.0"},
-		"sec-ch-ua-platform-version": {"19.0.0"},
-		"Sec-Fetch-Site":             {"same-origin"},
-		"Sec-Fetch-Mode":             {"cors"},
-		"Sec-Fetch-Dest":             {"empty"},
-		"Host":                       {"you.com"},
-	}
-
-	cookies := getCookies(dsToken)
-	var cookieStrings []string
-	for name, value := range cookies {
-		cookieStrings = append(cookieStrings, fmt.Sprintf("%s=%s", name, value))
-	}
-	youReq.Header.Add("Cookie", strings.Join(cookieStrings, ";"))
-
-	if !openAIReq.Stream {
-		handleNonStreamingResponse(w, youReq)
-		return
-	}
-
-	handleStreamingResponse(w, youReq)
+    // 新增模型列表接口
+    if r.URL.Path == "/hf/v1/models" {
+        w.Header().Set("Content-Type", "application/json")
+        w.Header().Set("Access-Control-Allow-Origin", "*")
+        w.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONS")
+        w.Header().Set("Access-Control-Allow-Headers", "*")
+
+        if r.Method == "OPTIONS" {
+            w.WriteHeader(http.StatusOK)
+            return
+        }
+
+        models := make([]ModelDetail, 0, len(modelMap))
+        created := time.Now().Unix()
+        for modelID := range modelMap {
+            models = append(models, ModelDetail{
+                ID:      modelID,
+                Object:  "model",
+                Created: created,
+                OwnedBy: "organization-owner",
+            })
+        }
+
+        response := ModelResponse{
+            Object: "list",
+            Data:   models,
+        }
+
+        json.NewEncoder(w).Encode(response)
+        return
+    }
+
+    // 原有逻辑保持不变
+    if r.URL.Path != "/hf/v1/chat/completions" {
+        w.Header().Set("Content-Type", "application/json")
+        json.NewEncoder(w).Encode(map[string]string{
+            "status":  "You2Api Service Running...",
+            "message": "MoLoveSze...",
+        })
+        return
+    }
+
+    w.Header().Set("Access-Control-Allow-Origin", "*")
+    w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+    w.Header().Set("Access-Control-Allow-Headers", "*")
+
+    if r.Method == "OPTIONS" {
+        w.WriteHeader(http.StatusOK)
+        return
+    }
+
+    authHeader := r.Header.Get("Authorization")
+    if !strings.HasPrefix(authHeader, "Bearer ") {
+        http.Error(w, "Missing or invalid authorization header", http.StatusUnauthorized)
+        return
+    }
+    dsToken := strings.TrimPrefix(authHeader, "Bearer ")
+
+    var openAIReq OpenAIRequest
+    if err := json.NewDecoder(r.Body).Decode(&openAIReq); err != nil {
+        http.Error(w, "Invalid request body", http.StatusBadRequest)
+        return
+    }
+
+    originalModel = openAIReq.Model
+    lastMessage := openAIReq.Messages[len(openAIReq.Messages)-1].Content
+    var chatHistory []map[string]interface{}
+    for _, msg := range openAIReq.Messages {
+        chatMsg := map[string]interface{}{
+            "question": msg.Content,
+            "answer":   "",
+        }
+        if msg.Role == "assistant" {
+            chatMsg["question"] = ""
+            chatMsg["answer"] = msg.Content
+        }
+        chatHistory = append(chatHistory, chatMsg)
+    }
+
+    chatHistoryJSON, _ := json.Marshal(chatHistory)
+
+    youReq, _ := http.NewRequest("GET", "https://you.com/api/streamingSearch", nil)
+
+    q := youReq.URL.Query()
+    q.Add("q", lastMessage)
+    q.Add("page", "1")
+    q.Add("count", "10")
+    q.Add("safeSearch", "Moderate")
+    q.Add("mkt", "zh-HK")
+    q.Add("enable_worklow_generation_ux", "true")
+    q.Add("domain", "youchat")
+    q.Add("use_personalization_extraction", "true")
+    q.Add("pastChatLength", fmt.Sprintf("%d", len(chatHistory)-1))
+    q.Add("selectedChatMode", "custom")
+    q.Add("selectedAiModel", mapModelName(openAIReq.Model))
+    q.Add("enable_agent_clarification_questions", "true")
+    q.Add("use_nested_youchat_updates", "true")
+    q.Add("chat", string(chatHistoryJSON))
+    youReq.URL.RawQuery = q.Encode()
+
+    youReq.Header = http.Header{
+        "sec-ch-ua-platform":         {"Windows"},
+        "Cache-Control":              {"no-cache"},
+        "sec-ch-ua":                  {`"Not(A:Brand";v="99", "Microsoft Edge";v="133", "Chromium";v="133"`},
+        "sec-ch-ua-bitness":          {"64"},
+        "sec-ch-ua-model":            {""},
+        "sec-ch-ua-mobile":           {"?0"},
+        "sec-ch-ua-arch":             {"x86"},
+        "sec-ch-ua-full-version":     {"133.0.3065.39"},
+        "Accept":                     {"text/event-stream"},
+        "User-Agent":                 {"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36 Edg/133.0.0.0"},
+        "sec-ch-ua-platform-version": {"19.0.0"},
+        "Sec-Fetch-Site":             {"same-origin"},
+        "Sec-Fetch-Mode":             {"cors"},
+        "Sec-Fetch-Dest":             {"empty"},
+        "Host":                       {"you.com"},
+    }
+
+    cookies := getCookies(dsToken)
+    var cookieStrings []string
+    for name, value := range cookies {
+        cookieStrings = append(cookieStrings, fmt.Sprintf("%s=%s", name, value))
+    }
+    youReq.Header.Add("Cookie", strings.Join(cookieStrings, ";"))
+
+    if !openAIReq.Stream {
+        handleNonStreamingResponse(w, youReq)
+        return
+    }
+
+    handleStreamingResponse(w, youReq)
 }
 
 func getCookies(dsToken string) map[string]string {
-	return map[string]string{
-		"guest_has_seen_legal_disclaimer": "true",
-		"youchat_personalization":         "true",
-		"DS":                              dsToken,
-		"you_subscription":                "youpro_standard_year",
-		"youpro_subscription":             "true",
-		"ai_model":                        "deepseek_r1",
-		"youchat_smart_learn":             "true",
-	}
+    return map[string]string{
+        "guest_has_seen_legal_disclaimer": "true",
+        "youchat_personalization":         "true",
+        "DS":                              dsToken,
+        "you_subscription":                "youpro_standard_year",
+        "youpro_subscription":             "true",
+        "ai_model":                        "deepseek_r1",
+        "youchat_smart_learn":             "true",
+    }
 }
 
 func handleNonStreamingResponse(w http.ResponseWriter, youReq *http.Request) {
-	client := &http.Client{
-		Timeout: 60 * time.Second,
-	}
-	resp, err := client.Do(youReq)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
-	defer resp.Body.Close()
-
-	var fullResponse strings.Builder
-	scanner := bufio.NewScanner(resp.Body)
-
-	buf := make([]byte, 0, 64*1024)
-	scanner.Buffer(buf, 1024*1024)
-
-	for scanner.Scan() {
-		line := scanner.Text()
-		if strings.HasPrefix(line, "event: youChatToken") {
-			scanner.Scan()
-			data := scanner.Text()
-			if !strings.HasPrefix(data, "data: ") {
-				continue
-			}
-			var token YouChatResponse
-			if err := json.Unmarshal([]byte(strings.TrimPrefix(data, "data: ")), &token); err != nil {
-				continue
-			}
-			fullResponse.WriteString(token.YouChatToken)
-		}
-	}
-
-	if scanner.Err() != nil {
-		http.Error(w, "Error reading response", http.StatusInternalServerError)
-		return
-	}
-
-	openAIResp := OpenAIResponse{
-		ID:      "chatcmpl-" + fmt.Sprintf("%d", time.Now().Unix()),
-		Object:  "chat.completion",
-		Created: time.Now().Unix(),
-		Model:   reverseMapModelName(mapModelName(originalModel)),
-		Choices: []OpenAIChoice{
-			{
-				Message: Message{
-					Role:    "assistant",
-					Content: fullResponse.String(),
-				},
-				Index:        0,
-				FinishReason: "stop",
-			},
-		},
-	}
-
-	w.Header().Set("Content-Type", "application/json")
-	if err := json.NewEncoder(w).Encode(openAIResp); err != nil {
-		http.Error(w, "Error encoding response", http.StatusInternalServerError)
-		return
-	}
+    client := &http.Client{
+        Timeout: 60 * time.Second,
+    }
+    resp, err := client.Do(youReq)
+    if err != nil {
+        http.Error(w, err.Error(), http.StatusInternalServerError)
+        return
+    }
+    defer resp.Body.Close()
+
+    var fullResponse strings.Builder
+    scanner := bufio.NewScanner(resp.Body)
+
+    buf := make([]byte, 0, 64*1024)
+    scanner.Buffer(buf, 1024*1024)
+
+    for scanner.Scan() {
+        line := scanner.Text()
+        if strings.HasPrefix(line, "event: youChatToken") {
+            scanner.Scan()
+            data := scanner.Text()
+            if !strings.HasPrefix(data, "data: ") {
+                continue
+            }
+            var token YouChatResponse
+            if err := json.Unmarshal([]byte(strings.TrimPrefix(data, "data: ")), &token); err != nil {
+                continue
+            }
+            fullResponse.WriteString(token.YouChatToken)
+        }
+    }
+
+    if scanner.Err() != nil {
+        http.Error(w, "Error reading response", http.StatusInternalServerError)
+        return
+    }
+
+    openAIResp := OpenAIResponse{
+        ID:      "chatcmpl-" + fmt.Sprintf("%d", time.Now().Unix()),
+        Object:  "chat.completion",
+        Created: time.Now().Unix(),
+        Model:   reverseMapModelName(mapModelName(originalModel)),
+        Choices: []OpenAIChoice{
+            {
+                Message: Message{
+                    Role:    "assistant",
+                    Content: fullResponse.String(),
+                },
+                Index:        0,
+                FinishReason: "stop",
+            },
+        },
+    }
+
+    w.Header().Set("Content-Type", "application/json")
+    if err := json.NewEncoder(w).Encode(openAIResp); err != nil {
+        http.Error(w, "Error encoding response", http.StatusInternalServerError)
+        return
+    }
 }
 
 func handleStreamingResponse(w http.ResponseWriter, youReq *http.Request) {
-	client := &http.Client{}
-	resp, err := client.Do(youReq)
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
-	defer resp.Body.Close()
-
-	w.Header().Set("Content-Type", "text/event-stream")
-	w.Header().Set("Cache-Control", "no-cache")
-	w.Header().Set("Connection", "keep-alive")
-
-	scanner := bufio.NewScanner(resp.Body)
-	for scanner.Scan() {
-		line := scanner.Text()
-
-		if strings.HasPrefix(line, "event: youChatToken") {
-			scanner.Scan()
-			data := scanner.Text()
-
-			var token YouChatResponse
-			json.Unmarshal([]byte(strings.TrimPrefix(data, "data: ")), &token)
-
-			openAIResp := OpenAIStreamResponse{
-				ID:      "chatcmpl-" + fmt.Sprintf("%d", time.Now().Unix()),
-				Object:  "chat.completion.chunk",
-				Created: time.Now().Unix(),
-				Model:   reverseMapModelName(mapModelName(originalModel)),
-				Choices: []Choice{
-					{
-						Delta: Delta{
-							Content: token.YouChatToken,
-						},
-						Index:        0,
-						FinishReason: "",
-					},
-				},
-			}
-
-			respBytes, _ := json.Marshal(openAIResp)
-			fmt.Fprintf(w, "data: %s\n\n", string(respBytes))
-			w.(http.Flusher).Flush()
-		}
-	}
+    client := &http.Client{}
+    resp, err := client.Do(youReq)
+    if err != nil {
+        http.Error(w, err.Error(), http.StatusInternalServerError)
+        return
+    }
+    defer resp.Body.Close()
+
+    w.Header().Set("Content-Type", "text/event-stream")
+    w.Header().Set("Cache-Control", "no-cache")
+    w.Header().Set("Connection", "keep-alive")
+
+    scanner := bufio.NewScanner(resp.Body)
+    for scanner.Scan() {
+        line := scanner.Text()
+
+        if strings.HasPrefix(line, "event: youChatToken") {
+            scanner.Scan()
+            data := scanner.Text()
+
+            var token YouChatResponse
+            json.Unmarshal([]byte(strings.TrimPrefix(data, "data: ")), &token)
+
+            openAIResp := OpenAIStreamResponse{
+                ID:      "chatcmpl-" + fmt.Sprintf("%d", time.Now().Unix()),
+                Object:  "chat.completion.chunk",
+                Created: time.Now().Unix(),
+                Model:   reverseMapModelName(mapModelName(originalModel)),
+                Choices: []Choice{
+                    {
+                        Delta: Delta{
+                            Content: token.YouChatToken,
+                        },
+                        Index:        0,
+                        FinishReason: "",
+                    },
+                },
+            }
+
+            respBytes, _ := json.Marshal(openAIResp)
+            fmt.Fprintf(w, "data: %s\n\n", string(respBytes))
+            w.(http.Flusher).Flush()
+        }
+    }
 }
 
 func main() {
-	http.HandleFunc("/", Handler)
-	fmt.Println("Server starting on :7860...")
-	if err := http.ListenAndServe(":7860", nil); err != nil {
-		fmt.Printf("Error starting server: %v\n", err)
-	}
+    http.HandleFunc("/", Handler)
+    fmt.Println("Server starting on :7860...")
+    if err := http.ListenAndServe(":7860", nil); err != nil {
+        fmt.Printf("Error starting server: %v\n", err)
+    }
 }