🎨 Support Baichuan2

providers/baichuan/base.go

@@ -0,0 +1,30 @@
+package baichuan
+
+import (
+	"one-api/providers/base"
+	"one-api/providers/openai"
+
+	"github.com/gin-gonic/gin"
+)
+
+// 定义供应商工厂
+type BaichuanProviderFactory struct{}
+
+// 创建 BaichuanProvider
+// https://platform.baichuan-ai.com/docs/api
+func (f BaichuanProviderFactory) Create(c *gin.Context) base.ProviderInterface {
+	return &BaichuanProvider{
+		OpenAIProvider: openai.OpenAIProvider{
+			BaseProvider: base.BaseProvider{
+				BaseURL:         "https://api.baichuan-ai.com",
+				ChatCompletions: "/v1/chat/completions",
+				Embeddings:      "/v1/embeddings",
+				Context:         c,
+			},
+		},
+	}
+}
+
+type BaichuanProvider struct {
+	openai.OpenAIProvider
+}

providers/baichuan/chat.go

@@ -0,0 +1,100 @@
+package baichuan
+
+import (
+	"net/http"
+	"one-api/common"
+	"one-api/providers/openai"
+	"one-api/types"
+	"strings"
+)
+
+func (baichuanResponse *BaichuanChatResponse) ResponseHandler(resp *http.Response) (OpenAIResponse any, errWithCode *types.OpenAIErrorWithStatusCode) {
+	if baichuanResponse.Error.Message != "" {
+		errWithCode = &types.OpenAIErrorWithStatusCode{
+			OpenAIError: baichuanResponse.Error,
+			StatusCode:  resp.StatusCode,
+		}
+
+		return
+	}
+
+	OpenAIResponse = types.ChatCompletionResponse{
+		ID:      baichuanResponse.ID,
+		Object:  baichuanResponse.Object,
+		Created: baichuanResponse.Created,
+		Model:   baichuanResponse.Model,
+		Choices: baichuanResponse.Choices,
+		Usage:   baichuanResponse.Usage,
+	}
+
+	return
+}
+
+// 获取聊天请求体
+func (p *BaichuanProvider) getChatRequestBody(request *types.ChatCompletionRequest) *BaichuanChatRequest {
+	messages := make([]BaichuanMessage, 0, len(request.Messages))
+	for i := 0; i < len(request.Messages); i++ {
+		message := request.Messages[i]
+		if message.Role == "system" || message.Role == "assistant" {
+			message.Role = "assistant"
+		} else {
+			message.Role = "user"
+		}
+		messages = append(messages, BaichuanMessage{
+			Content: message.StringContent(),
+			Role:    strings.ToLower(message.Role),
+		})
+	}
+
+	return &BaichuanChatRequest{
+		Model:       request.Model,
+		Messages:    messages,
+		Stream:      request.Stream,
+		Temperature: request.Temperature,
+		TopP:        request.TopP,
+		TopK:        request.N,
+	}
+}
+
+// 聊天
+func (p *BaichuanProvider) ChatAction(request *types.ChatCompletionRequest, isModelMapped bool, promptTokens int) (usage *types.Usage, errWithCode *types.OpenAIErrorWithStatusCode) {
+
+	requestBody := p.getChatRequestBody(request)
+
+	fullRequestURL := p.GetFullRequestURL(p.ChatCompletions, request.Model)
+	headers := p.GetRequestHeaders()
+	if request.Stream {
+		headers["Accept"] = "text/event-stream"
+	}
+
+	client := common.NewClient()
+	req, err := client.NewRequest(p.Context.Request.Method, fullRequestURL, common.WithBody(requestBody), common.WithHeader(headers))
+	if err != nil {
+		return nil, common.ErrorWrapper(err, "new_request_failed", http.StatusInternalServerError)
+	}
+
+	if request.Stream {
+		openAIProviderChatStreamResponse := &openai.OpenAIProviderChatStreamResponse{}
+		var textResponse string
+		errWithCode, textResponse = p.SendStreamRequest(req, openAIProviderChatStreamResponse)
+		if errWithCode != nil {
+			return
+		}
+
+		usage = &types.Usage{
+			PromptTokens:     promptTokens,
+			CompletionTokens: common.CountTokenText(textResponse, request.Model),
+			TotalTokens:      promptTokens + common.CountTokenText(textResponse, request.Model),
+		}
+
+	} else {
+		baichuanResponse := &BaichuanChatResponse{}
+		errWithCode = p.SendRequest(req, baichuanResponse, false)
+		if errWithCode != nil {
+			return
+		}
+
+		usage = baichuanResponse.Usage
+	}
+	return
+}

providers/baichuan/type.go

@@ -0,0 +1,36 @@
+package baichuan
+
+import "one-api/providers/openai"
+
+type BaichuanMessage struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type BaichuanKnowledgeBase struct {
+	Ids []string `json:"id"`
+}
+
+type BaichuanChatRequest struct {
+	Model             string                `json:"model"`
+	Messages          []BaichuanMessage     `json:"messages"`
+	Stream            bool                  `json:"stream,omitempty"`
+	Temperature       float64               `json:"temperature,omitempty"`
+	TopP              float64               `json:"top_p,omitempty"`
+	TopK              int                   `json:"top_k,omitempty"`
+	WithSearchEnhance bool                  `json:"with_search_enhance,omitempty"`
+	KnowledgeBase     BaichuanKnowledgeBase `json:"knowledge_base,omitempty"`
+}
+
+type BaichuanKnowledgeBaseResponse struct {
+	Cites []struct {
+		Title   string `json:"title"`
+		Content string `json:"content"`
+		FileId  string `json:"file_id"`
+	} `json:"cites"`
+}
+
+type BaichuanChatResponse struct {
+	openai.OpenAIProviderChatResponse
+	KnowledgeBase BaichuanKnowledgeBaseResponse `json:"knowledge_base,omitempty"`
+}