288 lines
9.1 KiB
Go
288 lines
9.1 KiB
Go
package claude
|
||
|
||
import (
|
||
"bufio"
|
||
"encoding/json"
|
||
"fmt"
|
||
"github.com/gin-gonic/gin"
|
||
"io"
|
||
"net/http"
|
||
"one-api/common"
|
||
"one-api/dto"
|
||
"one-api/service"
|
||
"strings"
|
||
)
|
||
|
||
func stopReasonClaude2OpenAI(reason string) string {
|
||
switch reason {
|
||
case "stop_sequence":
|
||
return "stop"
|
||
case "max_tokens":
|
||
return "length"
|
||
default:
|
||
return reason
|
||
}
|
||
}
|
||
|
||
func requestOpenAI2ClaudeComplete(textRequest dto.GeneralOpenAIRequest) *ClaudeRequest {
|
||
claudeRequest := ClaudeRequest{
|
||
Model: textRequest.Model,
|
||
Prompt: "",
|
||
MaxTokensToSample: textRequest.MaxTokens,
|
||
StopSequences: nil,
|
||
Temperature: textRequest.Temperature,
|
||
TopP: textRequest.TopP,
|
||
Stream: textRequest.Stream,
|
||
}
|
||
if claudeRequest.MaxTokensToSample == 0 {
|
||
claudeRequest.MaxTokensToSample = 1000000
|
||
}
|
||
prompt := ""
|
||
for _, message := range textRequest.Messages {
|
||
if message.Role == "user" {
|
||
prompt += fmt.Sprintf("\n\nHuman: %s", message.Content)
|
||
} else if message.Role == "assistant" {
|
||
prompt += fmt.Sprintf("\n\nAssistant: %s", message.Content)
|
||
} else if message.Role == "system" {
|
||
if prompt == "" {
|
||
prompt = message.StringContent()
|
||
}
|
||
}
|
||
}
|
||
prompt += "\n\nAssistant:"
|
||
claudeRequest.Prompt = prompt
|
||
return &claudeRequest
|
||
}
|
||
|
||
func requestOpenAI2ClaudeMessage(textRequest dto.GeneralOpenAIRequest) (*ClaudeRequest, error) {
|
||
claudeRequest := ClaudeRequest{
|
||
Model: textRequest.Model,
|
||
MaxTokens: textRequest.MaxTokens,
|
||
StopSequences: nil,
|
||
Temperature: textRequest.Temperature,
|
||
TopP: textRequest.TopP,
|
||
Stream: textRequest.Stream,
|
||
}
|
||
claudeMessages := make([]ClaudeMessage, 0)
|
||
for _, message := range textRequest.Messages {
|
||
if message.Role == "system" {
|
||
claudeRequest.System = message.StringContent()
|
||
} else {
|
||
claudeMessage := ClaudeMessage{
|
||
Role: message.Role,
|
||
}
|
||
if message.IsStringContent() {
|
||
claudeMessage.Content = message.StringContent()
|
||
} else {
|
||
claudeMediaMessages := make([]ClaudeMediaMessage, 0)
|
||
for _, mediaMessage := range message.ParseContent() {
|
||
claudeMediaMessage := ClaudeMediaMessage{
|
||
Type: mediaMessage.Type,
|
||
}
|
||
if mediaMessage.Type == "text" {
|
||
claudeMediaMessage.Text = mediaMessage.Text
|
||
} else {
|
||
imageUrl := mediaMessage.ImageUrl.(dto.MessageImageUrl)
|
||
claudeMediaMessage.Type = "image"
|
||
claudeMediaMessage.Source = &ClaudeMessageSource{
|
||
Type: "base64",
|
||
}
|
||
// 判断是否是url
|
||
if strings.HasPrefix(imageUrl.Url, "http") {
|
||
// 是url，获取图片的类型和base64编码的数据
|
||
mimeType, data, _ := common.GetImageFromUrl(imageUrl.Url)
|
||
claudeMediaMessage.Source.MediaType = mimeType
|
||
claudeMediaMessage.Source.Data = data
|
||
} else {
|
||
_, format, base64String, err := common.DecodeBase64ImageData(imageUrl.Url)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
claudeMediaMessage.Source.MediaType = "image/" + format
|
||
claudeMediaMessage.Source.Data = base64String
|
||
}
|
||
}
|
||
claudeMediaMessages = append(claudeMediaMessages, claudeMediaMessage)
|
||
}
|
||
claudeMessage.Content = claudeMediaMessages
|
||
}
|
||
claudeMessages = append(claudeMessages, claudeMessage)
|
||
}
|
||
}
|
||
claudeRequest.Prompt = ""
|
||
claudeRequest.Messages = claudeMessages
|
||
reqJson, _ := json.Marshal(claudeRequest)
|
||
common.SysLog(fmt.Sprintf("claude request: %s", reqJson))
|
||
|
||
return &claudeRequest, nil
|
||
}
|
||
|
||
func streamResponseClaude2OpenAI(claudeResponse *ClaudeResponse) *dto.ChatCompletionsStreamResponse {
|
||
var choice dto.ChatCompletionsStreamResponseChoice
|
||
choice.Delta.Content = claudeResponse.Completion
|
||
finishReason := stopReasonClaude2OpenAI(claudeResponse.StopReason)
|
||
if finishReason != "null" {
|
||
choice.FinishReason = &finishReason
|
||
}
|
||
var response dto.ChatCompletionsStreamResponse
|
||
response.Object = "chat.completion.chunk"
|
||
response.Model = claudeResponse.Model
|
||
response.Choices = []dto.ChatCompletionsStreamResponseChoice{choice}
|
||
return &response
|
||
}
|
||
|
||
func responseClaude2OpenAI(reqMode int, claudeResponse *ClaudeResponse) *dto.OpenAITextResponse {
|
||
choices := make([]dto.OpenAITextResponseChoice, 0)
|
||
fullTextResponse := dto.OpenAITextResponse{
|
||
Id: fmt.Sprintf("chatcmpl-%s", common.GetUUID()),
|
||
Object: "chat.completion",
|
||
Created: common.GetTimestamp(),
|
||
}
|
||
if reqMode == RequestModeCompletion {
|
||
content, _ := json.Marshal(strings.TrimPrefix(claudeResponse.Completion, " "))
|
||
choice := dto.OpenAITextResponseChoice{
|
||
Index: 0,
|
||
Message: dto.Message{
|
||
Role: "assistant",
|
||
Content: content,
|
||
Name: nil,
|
||
},
|
||
FinishReason: stopReasonClaude2OpenAI(claudeResponse.StopReason),
|
||
}
|
||
choices = append(choices, choice)
|
||
} else {
|
||
fullTextResponse.Id = claudeResponse.Id
|
||
for i, message := range claudeResponse.Content {
|
||
content, _ := json.Marshal(message.Text)
|
||
choice := dto.OpenAITextResponseChoice{
|
||
Index: i,
|
||
Message: dto.Message{
|
||
Role: "assistant",
|
||
Content: content,
|
||
},
|
||
FinishReason: stopReasonClaude2OpenAI(claudeResponse.StopReason),
|
||
}
|
||
choices = append(choices, choice)
|
||
}
|
||
}
|
||
|
||
fullTextResponse.Choices = choices
|
||
return &fullTextResponse
|
||
}
|
||
|
||
func claudeStreamHandler(c *gin.Context, resp *http.Response) (*dto.OpenAIErrorWithStatusCode, string) {
|
||
responseText := ""
|
||
responseId := fmt.Sprintf("chatcmpl-%s", common.GetUUID())
|
||
createdTime := common.GetTimestamp()
|
||
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 := strings.Index(string(data), "\r\n\r\n"); i >= 0 {
|
||
return i + 4, 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 !strings.HasPrefix(data, "event: completion") {
|
||
continue
|
||
}
|
||
data = strings.TrimPrefix(data, "event: completion\r\ndata: ")
|
||
dataChan <- data
|
||
}
|
||
stopChan <- true
|
||
}()
|
||
service.SetEventStreamHeaders(c)
|
||
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 claudeResponse ClaudeResponse
|
||
err := json.Unmarshal([]byte(data), &claudeResponse)
|
||
if err != nil {
|
||
common.SysError("error unmarshalling stream response: " + err.Error())
|
||
return true
|
||
}
|
||
responseText += claudeResponse.Completion
|
||
response := streamResponseClaude2OpenAI(&claudeResponse)
|
||
response.Id = responseId
|
||
response.Created = createdTime
|
||
jsonStr, err := json.Marshal(response)
|
||
if err != nil {
|
||
common.SysError("error marshalling stream response: " + err.Error())
|
||
return true
|
||
}
|
||
c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonStr)})
|
||
return true
|
||
case <-stopChan:
|
||
c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
|
||
return false
|
||
}
|
||
})
|
||
err := resp.Body.Close()
|
||
if err != nil {
|
||
return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), ""
|
||
}
|
||
return nil, responseText
|
||
}
|
||
|
||
func claudeHandler(requestMode int, c *gin.Context, resp *http.Response, promptTokens int, model string) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) {
|
||
responseBody, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil
|
||
}
|
||
err = resp.Body.Close()
|
||
if err != nil {
|
||
return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
|
||
}
|
||
var claudeResponse ClaudeResponse
|
||
common.SysLog(fmt.Sprintf("claude response: %s", responseBody))
|
||
err = json.Unmarshal(responseBody, &claudeResponse)
|
||
if err != nil {
|
||
return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
|
||
}
|
||
respJson, _ := json.Marshal(claudeResponse)
|
||
common.SysLog(fmt.Sprintf("claude response json: %s", respJson))
|
||
if claudeResponse.Error.Type != "" {
|
||
return &dto.OpenAIErrorWithStatusCode{
|
||
Error: dto.OpenAIError{
|
||
Message: claudeResponse.Error.Message,
|
||
Type: claudeResponse.Error.Type,
|
||
Param: "",
|
||
Code: claudeResponse.Error.Type,
|
||
},
|
||
StatusCode: resp.StatusCode,
|
||
}, nil
|
||
}
|
||
fullTextResponse := responseClaude2OpenAI(requestMode, &claudeResponse)
|
||
completionTokens := service.CountTokenText(claudeResponse.Completion, model)
|
||
usage := dto.Usage{}
|
||
if requestMode == RequestModeCompletion {
|
||
usage.PromptTokens = promptTokens
|
||
usage.CompletionTokens = completionTokens
|
||
usage.TotalTokens = promptTokens + completionTokens
|
||
} else {
|
||
usage.PromptTokens = claudeResponse.Usage.InputTokens
|
||
usage.CompletionTokens = claudeResponse.Usage.OutputTokens
|
||
usage.TotalTokens = claudeResponse.Usage.InputTokens + claudeResponse.Usage.OutputTokens
|
||
}
|
||
fullTextResponse.Usage = usage
|
||
jsonResponse, err := json.Marshal(fullTextResponse)
|
||
if err != nil {
|
||
return service.OpenAIErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil
|
||
}
|
||
c.Writer.Header().Set("Content-Type", "application/json")
|
||
c.Writer.WriteHeader(resp.StatusCode)
|
||
_, err = c.Writer.Write(jsonResponse)
|
||
return nil, &usage
|
||
}
|