280 lines
7.9 KiB
Go
280 lines
7.9 KiB
Go
// https://docs.anthropic.com/claude/reference/messages_post
|
|
|
|
package claude
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"opencatd-open/pkg/openai"
|
|
"opencatd-open/pkg/tokenizer"
|
|
"opencatd-open/store"
|
|
"strings"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func ChatProxy(c *gin.Context, chatReq *openai.ChatCompletionRequest) {
|
|
ChatMessages(c, chatReq)
|
|
}
|
|
|
|
func ChatTextCompletions(c *gin.Context, chatReq *openai.ChatCompletionRequest) {
|
|
|
|
}
|
|
|
|
type ChatRequest struct {
|
|
Model string `json:"model,omitempty"`
|
|
Messages any `json:"messages,omitempty"`
|
|
MaxTokens int `json:"max_tokens,omitempty"`
|
|
Stream bool `json:"stream,omitempty"`
|
|
System string `json:"system,omitempty"`
|
|
TopK int `json:"top_k,omitempty"`
|
|
TopP float64 `json:"top_p,omitempty"`
|
|
Temperature float64 `json:"temperature,omitempty"`
|
|
}
|
|
|
|
func (c *ChatRequest) ByteJson() []byte {
|
|
bytejson, _ := json.Marshal(c)
|
|
return bytejson
|
|
}
|
|
|
|
type ChatMessage struct {
|
|
Role string `json:"role,omitempty"`
|
|
Content string `json:"content,omitempty"`
|
|
}
|
|
|
|
type VisionMessages struct {
|
|
Role string `json:"role,omitempty"`
|
|
Content []VisionContent `json:"content,omitempty"`
|
|
}
|
|
|
|
type VisionContent struct {
|
|
Type string `json:"type,omitempty"`
|
|
Source *VisionSource `json:"source,omitempty"`
|
|
Text string `json:"text,omitempty"`
|
|
}
|
|
|
|
type VisionSource struct {
|
|
Type string `json:"type,omitempty"`
|
|
MediaType string `json:"media_type,omitempty"`
|
|
Data string `json:"data,omitempty"`
|
|
}
|
|
|
|
type ChatResponse struct {
|
|
ID string `json:"id"`
|
|
Type string `json:"type"`
|
|
Role string `json:"role"`
|
|
Model string `json:"model"`
|
|
StopSequence any `json:"stop_sequence"`
|
|
Usage struct {
|
|
InputTokens int `json:"input_tokens"`
|
|
OutputTokens int `json:"output_tokens"`
|
|
} `json:"usage"`
|
|
Content []struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text"`
|
|
} `json:"content"`
|
|
StopReason string `json:"stop_reason"`
|
|
}
|
|
|
|
type ClaudeStreamResponse struct {
|
|
Type string `json:"type"`
|
|
Index int `json:"index"`
|
|
ContentBlock struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text"`
|
|
} `json:"content_block"`
|
|
Delta struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text"`
|
|
StopReason string `json:"stop_reason"`
|
|
StopSequence any `json:"stop_sequence"`
|
|
} `json:"delta"`
|
|
Message struct {
|
|
ID string `json:"id"`
|
|
Type string `json:"type"`
|
|
Role string `json:"role"`
|
|
Content []any `json:"content"`
|
|
Model string `json:"model"`
|
|
StopReason string `json:"stop_reason"`
|
|
StopSequence any `json:"stop_sequence"`
|
|
Usage struct {
|
|
InputTokens int `json:"input_tokens"`
|
|
OutputTokens int `json:"output_tokens"`
|
|
} `json:"usage"`
|
|
} `json:"message"`
|
|
Error struct {
|
|
Type string `json:"type"`
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
Usage struct {
|
|
OutputTokens int `json:"output_tokens"`
|
|
} `json:"usage"`
|
|
}
|
|
|
|
func ChatMessages(c *gin.Context, chatReq *openai.ChatCompletionRequest) {
|
|
// var haveImages bool
|
|
|
|
usagelog := store.Tokens{Model: chatReq.Model}
|
|
var claudReq ChatRequest
|
|
claudReq.Model = chatReq.Model
|
|
claudReq.Stream = chatReq.Stream
|
|
claudReq.Temperature = chatReq.Temperature
|
|
claudReq.TopP = chatReq.TopP
|
|
claudReq.MaxTokens = 4096
|
|
|
|
var prompt string
|
|
|
|
var claudecontent []VisionContent
|
|
for _, msg := range chatReq.Messages {
|
|
if msg.Role == "system" {
|
|
claudReq.System = string(msg.Content)
|
|
continue
|
|
}
|
|
|
|
var oaivisioncontent []openai.VisionContent
|
|
if err := json.Unmarshal(msg.Content, &oaivisioncontent); err != nil {
|
|
prompt += "<" + msg.Role + ">: " + string(msg.Content) + "\n"
|
|
|
|
claudecontent = append(claudecontent, VisionContent{Type: "text", Text: msg.Role + ":" + string(msg.Content)})
|
|
} else {
|
|
if len(oaivisioncontent) > 0 {
|
|
for _, content := range oaivisioncontent {
|
|
if content.Type == "text" {
|
|
prompt += "<" + msg.Role + ">: " + content.Text + "\n"
|
|
claudecontent = append(claudecontent, VisionContent{Type: "text", Text: msg.Role + ":" + content.Text})
|
|
} else if content.Type == "image_url" {
|
|
if strings.HasPrefix(content.ImageURL.URL, "http") {
|
|
fmt.Println("链接:", content.ImageURL.URL)
|
|
} else if strings.HasPrefix(content.ImageURL.URL, "data:image") {
|
|
fmt.Println("base64:", content.ImageURL.URL[:20])
|
|
}
|
|
// todo image tokens
|
|
var mediaType string
|
|
if strings.HasPrefix(content.ImageURL.URL, "data:image/jpeg") {
|
|
mediaType = "image/jpeg"
|
|
}
|
|
if strings.HasPrefix(content.ImageURL.URL, "data:image/png") {
|
|
mediaType = "image/png"
|
|
}
|
|
claudecontent = append(claudecontent, VisionContent{Type: "image", Source: &VisionSource{Type: "base64", MediaType: mediaType, Data: strings.Split(content.ImageURL.URL, ",")[1]}})
|
|
}
|
|
}
|
|
|
|
}
|
|
}
|
|
// if len(chatReq.Tools) > 0 {
|
|
// tooljson, _ := json.Marshal(chatReq.Tools)
|
|
// prompt += "<tools>: " + string(tooljson) + "\n"
|
|
// }
|
|
}
|
|
claudReq.Messages = []VisionMessages{{Role: "user", Content: claudecontent}}
|
|
|
|
usagelog.PromptCount = tokenizer.NumTokensFromStr(prompt, chatReq.Model)
|
|
|
|
req, _ := http.NewRequest("POST", MessageEndpoint, strings.NewReader(fmt.Sprintf("%v", bytes.NewReader(claudReq.ByteJson()))))
|
|
client := http.DefaultClient
|
|
rsp, err := client.Do(req)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
defer rsp.Body.Close()
|
|
if rsp.StatusCode != http.StatusOK {
|
|
io.Copy(c.Writer, rsp.Body)
|
|
return
|
|
}
|
|
|
|
teeReader := io.TeeReader(rsp.Body, c.Writer)
|
|
|
|
dataChan := make(chan string, 1)
|
|
// stopChan := make(chan bool)
|
|
|
|
var result string
|
|
|
|
scanner := bufio.NewScanner(teeReader)
|
|
|
|
go func() {
|
|
for scanner.Scan() {
|
|
line := scanner.Bytes()
|
|
if len(line) > 0 && bytes.HasPrefix(line, []byte("data: ")) {
|
|
if bytes.HasPrefix(line, []byte("data: [DONE]")) {
|
|
dataChan <- string(line) + "\n"
|
|
break
|
|
}
|
|
var claudeResp ClaudeStreamResponse
|
|
line = bytes.Replace(line, []byte("data: "), []byte(""), -1)
|
|
line = bytes.TrimSpace(line)
|
|
if err := json.Unmarshal(line, &claudeResp); err != nil {
|
|
continue
|
|
}
|
|
|
|
if claudeResp.Type == "message_start" {
|
|
if claudeResp.Message.Role != "" {
|
|
result += "<" + claudeResp.Message.Role + ">"
|
|
}
|
|
} else if claudeResp.Type == "message_stop" {
|
|
break
|
|
}
|
|
|
|
if claudeResp.Delta.Text != "" {
|
|
result += claudeResp.Delta.Text
|
|
}
|
|
var choice openai.Choice
|
|
choice.Delta.Role = claudeResp.Message.Role
|
|
choice.Delta.Content = claudeResp.Delta.Text
|
|
choice.FinishReason = claudeResp.Delta.StopReason
|
|
|
|
chatResp := openai.ChatCompletionStreamResponse{
|
|
Model: chatReq.Model,
|
|
Choices: []openai.Choice{choice},
|
|
}
|
|
dataChan <- "data: " + string(chatResp.ByteJson()) + "\n"
|
|
if claudeResp.Delta.StopReason != "" {
|
|
dataChan <- "\ndata: [DONE]\n"
|
|
}
|
|
} else {
|
|
if !bytes.HasPrefix(line, []byte("event:")) {
|
|
dataChan <- string(line) + "\n"
|
|
}
|
|
}
|
|
}
|
|
defer close(dataChan)
|
|
}()
|
|
|
|
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
|
c.Writer.Header().Set("Cache-Control", "no-cache")
|
|
c.Writer.Header().Set("Connection", "keep-alive")
|
|
c.Writer.Header().Set("Transfer-Encoding", "chunked")
|
|
c.Writer.Header().Set("X-Accel-Buffering", "no")
|
|
|
|
c.Stream(func(w io.Writer) bool {
|
|
if data, ok := <-dataChan; ok {
|
|
if strings.HasPrefix(data, "data: ") {
|
|
c.Writer.WriteString(data)
|
|
// c.Writer.WriteString("\n\n")
|
|
} else {
|
|
c.Writer.WriteHeader(http.StatusBadGateway)
|
|
c.Writer.WriteString(data)
|
|
}
|
|
c.Writer.Flush()
|
|
return true
|
|
}
|
|
go func() {
|
|
usagelog.CompletionCount = tokenizer.NumTokensFromStr(result, chatReq.Model)
|
|
usagelog.Cost = fmt.Sprintf("%.6f", tokenizer.Cost(usagelog.Model, usagelog.PromptCount, usagelog.CompletionCount))
|
|
if err := store.Record(&usagelog); err != nil {
|
|
log.Println(err)
|
|
}
|
|
if err := store.SumDaily(usagelog.UserID); err != nil {
|
|
log.Println(err)
|
|
}
|
|
}()
|
|
return false
|
|
})
|
|
}
|