Rain
302 lines
9.0 KiB
Go
302 lines
9.0 KiB
Go
package dify
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"mime/multipart"
|
||
"net/http"
|
||
"os"
|
||
"strings"
|
||
|
||
"github.com/QuantumNous/new-api/common"
|
||
"github.com/QuantumNous/new-api/constant"
|
||
"github.com/QuantumNous/new-api/dto"
|
||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||
"github.com/QuantumNous/new-api/relay/helper"
|
||
"github.com/QuantumNous/new-api/service"
|
||
"github.com/QuantumNous/new-api/types"
|
||
"github.com/samber/lo"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
func uploadDifyFile(c *gin.Context, info *relaycommon.RelayInfo, user string, media dto.MediaContent) *DifyFile {
|
||
uploadUrl := fmt.Sprintf("%s/v1/files/upload", info.ChannelBaseUrl)
|
||
switch media.Type {
|
||
case dto.ContentTypeImageURL:
|
||
// Decode base64 data
|
||
imageMedia := media.GetImageMedia()
|
||
base64Data := imageMedia.Url
|
||
// Remove base64 prefix if exists (e.g., "data:image/jpeg;base64,")
|
||
if idx := strings.Index(base64Data, ","); idx != -1 {
|
||
base64Data = base64Data[idx+1:]
|
||
}
|
||
|
||
// Decode base64 string
|
||
decodedData, err := base64.StdEncoding.DecodeString(base64Data)
|
||
if err != nil {
|
||
common.SysLog("failed to decode base64: " + err.Error())
|
||
return nil
|
||
}
|
||
|
||
// Create temporary file
|
||
tempFile, err := os.CreateTemp("", "dify-upload-*")
|
||
if err != nil {
|
||
common.SysLog("failed to create temp file: " + err.Error())
|
||
return nil
|
||
}
|
||
defer tempFile.Close()
|
||
defer os.Remove(tempFile.Name())
|
||
|
||
// Write decoded data to temp file
|
||
if _, err := tempFile.Write(decodedData); err != nil {
|
||
common.SysLog("failed to write to temp file: " + err.Error())
|
||
return nil
|
||
}
|
||
|
||
// Create multipart form
|
||
body := &bytes.Buffer{}
|
||
writer := multipart.NewWriter(body)
|
||
|
||
// Add user field
|
||
if err := writer.WriteField("user", user); err != nil {
|
||
common.SysLog("failed to add user field: " + err.Error())
|
||
return nil
|
||
}
|
||
|
||
// Create form file with proper mime type
|
||
mimeType := imageMedia.MimeType
|
||
if mimeType == "" {
|
||
mimeType = "image/jpeg" // default mime type
|
||
}
|
||
|
||
// Create form file
|
||
part, err := writer.CreateFormFile("file", fmt.Sprintf("image.%s", strings.TrimPrefix(mimeType, "image/")))
|
||
if err != nil {
|
||
common.SysLog("failed to create form file: " + err.Error())
|
||
return nil
|
||
}
|
||
|
||
// Copy file content to form
|
||
if _, err = io.Copy(part, bytes.NewReader(decodedData)); err != nil {
|
||
common.SysLog("failed to copy file content: " + err.Error())
|
||
return nil
|
||
}
|
||
writer.Close()
|
||
|
||
// Create HTTP request
|
||
req, err := http.NewRequest("POST", uploadUrl, body)
|
||
if err != nil {
|
||
common.SysLog("failed to create request: " + err.Error())
|
||
return nil
|
||
}
|
||
|
||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
|
||
|
||
// Send request
|
||
client := service.GetHttpClient()
|
||
resp, err := client.Do(req)
|
||
if err != nil {
|
||
common.SysLog("failed to send request: " + err.Error())
|
||
return nil
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
// Parse response
|
||
var result struct {
|
||
Id string `json:"id"`
|
||
}
|
||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||
common.SysLog("failed to decode response: " + err.Error())
|
||
return nil
|
||
}
|
||
|
||
return &DifyFile{
|
||
UploadFileId: result.Id,
|
||
Type: "image",
|
||
TransferMode: "local_file",
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func requestOpenAI2Dify(c *gin.Context, info *relaycommon.RelayInfo, request dto.GeneralOpenAIRequest) *DifyChatRequest {
|
||
difyReq := DifyChatRequest{
|
||
Inputs: make(map[string]interface{}),
|
||
AutoGenerateName: false,
|
||
}
|
||
|
||
user := request.User
|
||
if len(user) == 0 {
|
||
user = json.RawMessage(helper.GetResponseID(c))
|
||
}
|
||
var stringUser string
|
||
err := json.Unmarshal(user, &stringUser)
|
||
if err != nil {
|
||
common.SysLog("failed to unmarshal user: " + err.Error())
|
||
stringUser = helper.GetResponseID(c)
|
||
}
|
||
difyReq.User = stringUser
|
||
|
||
files := make([]DifyFile, 0)
|
||
var content strings.Builder
|
||
for _, message := range request.Messages {
|
||
if message.Role == "system" {
|
||
content.WriteString("SYSTEM: \n" + message.StringContent() + "\n")
|
||
} else if message.Role == "assistant" {
|
||
content.WriteString("ASSISTANT: \n" + message.StringContent() + "\n")
|
||
} else {
|
||
parseContent := message.ParseContent()
|
||
for _, mediaContent := range parseContent {
|
||
switch mediaContent.Type {
|
||
case dto.ContentTypeText:
|
||
content.WriteString("USER: \n" + mediaContent.Text + "\n")
|
||
case dto.ContentTypeImageURL:
|
||
media := mediaContent.GetImageMedia()
|
||
var file *DifyFile
|
||
if media.IsRemoteImage() {
|
||
// 修复 #2083: 远程图片分支此前未初始化 file，
|
||
// 导致 file.Type = ... 触发 nil pointer dereference
|
||
// 而 panic（500: "invalid memory address or nil pointer dereference"）。
|
||
file = &DifyFile{
|
||
Type: media.MimeType,
|
||
TransferMode: "remote_url",
|
||
URL: media.Url,
|
||
}
|
||
} else {
|
||
file = uploadDifyFile(c, info, difyReq.User, mediaContent)
|
||
}
|
||
if file != nil {
|
||
files = append(files, *file)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
difyReq.Query = content.String()
|
||
difyReq.Files = files
|
||
mode := "blocking"
|
||
if lo.FromPtrOr(request.Stream, false) {
|
||
mode = "streaming"
|
||
}
|
||
difyReq.ResponseMode = mode
|
||
return &difyReq
|
||
}
|
||
|
||
func streamResponseDify2OpenAI(difyResponse DifyChunkChatCompletionResponse) *dto.ChatCompletionsStreamResponse {
|
||
response := dto.ChatCompletionsStreamResponse{
|
||
Object: "chat.completion.chunk",
|
||
Created: common.GetTimestamp(),
|
||
Model: "dify",
|
||
}
|
||
var choice dto.ChatCompletionsStreamResponseChoice
|
||
if strings.HasPrefix(difyResponse.Event, "workflow_") {
|
||
if constant.DifyDebug {
|
||
text := "Workflow: " + difyResponse.Data.WorkflowId
|
||
if difyResponse.Event == "workflow_finished" {
|
||
text += " " + difyResponse.Data.Status
|
||
}
|
||
choice.Delta.SetReasoningContent(text + "\n")
|
||
}
|
||
} else if strings.HasPrefix(difyResponse.Event, "node_") {
|
||
if constant.DifyDebug {
|
||
text := "Node: " + difyResponse.Data.NodeType
|
||
if difyResponse.Event == "node_finished" {
|
||
text += " " + difyResponse.Data.Status
|
||
}
|
||
choice.Delta.SetReasoningContent(text + "\n")
|
||
}
|
||
} else if difyResponse.Event == "message" || difyResponse.Event == "agent_message" {
|
||
if difyResponse.Answer == "<details style=\"color:gray;background-color: #f8f8f8;padding: 8px;border-radius: 4px;\" open> <summary> Thinking... </summary>\n" {
|
||
difyResponse.Answer = "<think>"
|
||
} else if difyResponse.Answer == "</details>" {
|
||
difyResponse.Answer = "</think>"
|
||
}
|
||
|
||
choice.Delta.SetContentString(difyResponse.Answer)
|
||
}
|
||
response.Choices = append(response.Choices, choice)
|
||
return &response
|
||
}
|
||
|
||
func difyStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||
var responseText string
|
||
usage := &dto.Usage{}
|
||
var nodeToken int
|
||
helper.SetEventStreamHeaders(c)
|
||
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
|
||
var difyResponse DifyChunkChatCompletionResponse
|
||
if err := json.Unmarshal([]byte(data), &difyResponse); err != nil {
|
||
common.SysLog("error unmarshalling stream response: " + err.Error())
|
||
sr.Error(err)
|
||
return
|
||
}
|
||
if difyResponse.Event == "message_end" {
|
||
usage = &difyResponse.MetaData.Usage
|
||
sr.Done()
|
||
return
|
||
} else if difyResponse.Event == "error" {
|
||
sr.Stop(fmt.Errorf("dify error event"))
|
||
return
|
||
}
|
||
openaiResponse := *streamResponseDify2OpenAI(difyResponse)
|
||
if len(openaiResponse.Choices) != 0 {
|
||
responseText += openaiResponse.Choices[0].Delta.GetContentString()
|
||
if openaiResponse.Choices[0].Delta.ReasoningContent != nil {
|
||
nodeToken += 1
|
||
}
|
||
}
|
||
if err := helper.ObjectData(c, openaiResponse); err != nil {
|
||
common.SysLog(err.Error())
|
||
sr.Error(err)
|
||
}
|
||
})
|
||
helper.Done(c)
|
||
if usage.TotalTokens == 0 {
|
||
usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens())
|
||
}
|
||
usage.CompletionTokens += nodeToken
|
||
return usage, nil
|
||
}
|
||
|
||
func difyHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
|
||
var difyResponse DifyChatCompletionResponse
|
||
responseBody, err := io.ReadAll(resp.Body)
|
||
|
||
if err != nil {
|
||
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
|
||
}
|
||
service.CloseResponseBodyGracefully(resp)
|
||
err = json.Unmarshal(responseBody, &difyResponse)
|
||
if err != nil {
|
||
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
|
||
}
|
||
fullTextResponse := dto.OpenAITextResponse{
|
||
Id: difyResponse.ConversationId,
|
||
Object: "chat.completion",
|
||
Created: common.GetTimestamp(),
|
||
Usage: difyResponse.MetaData.Usage,
|
||
}
|
||
choice := dto.OpenAITextResponseChoice{
|
||
Index: 0,
|
||
Message: dto.Message{
|
||
Role: "assistant",
|
||
Content: difyResponse.Answer,
|
||
},
|
||
FinishReason: "stop",
|
||
}
|
||
fullTextResponse.Choices = append(fullTextResponse.Choices, choice)
|
||
jsonResponse, err := json.Marshal(fullTextResponse)
|
||
if err != nil {
|
||
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
|
||
}
|
||
c.Writer.Header().Set("Content-Type", "application/json")
|
||
c.Writer.WriteHeader(resp.StatusCode)
|
||
c.Writer.Write(jsonResponse)
|
||
return &difyResponse.MetaData.Usage, nil
|
||
}
|