reface to openteam
This commit is contained in:
Sakurasan
244
llm/google/chat.go
Normal file
244
llm/google/chat.go
Normal file
@@ -0,0 +1,244 @@
|
||||
// https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/grounding-search-entry-points?authuser=2&hl=zh-cn
|
||||
//
|
||||
// https://cloud.google.com/vertex-ai/docs/generative-ai/quotas-genai
|
||||
package google
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"opencatd-open/llm/openai"
|
||||
"opencatd-open/pkg/tokenizer"
|
||||
"opencatd-open/store"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/generative-ai-go/genai"
|
||||
"google.golang.org/api/iterator"
|
||||
"google.golang.org/api/option"
|
||||
)
|
||||
|
||||
type GeminiChatRequest struct {
|
||||
Contents []GeminiContent `json:"contents,omitempty"`
|
||||
}
|
||||
|
||||
func (g GeminiChatRequest) ByteJson() []byte {
|
||||
bytejson, _ := json.Marshal(g)
|
||||
return bytejson
|
||||
}
|
||||
|
||||
type GeminiContent struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Parts []GeminiPart `json:"parts,omitempty"`
|
||||
}
|
||||
type GeminiPart struct {
|
||||
Text string `json:"text,omitempty"`
|
||||
// InlineData GeminiPartInlineData `json:"inlineData,omitempty"`
|
||||
}
|
||||
type GeminiPartInlineData struct {
|
||||
MimeType string `json:"mimeType,omitempty"`
|
||||
Data string `json:"data,omitempty"` // base64
|
||||
}
|
||||
|
||||
type GeminiResponse struct {
|
||||
Candidates []struct {
|
||||
Content struct {
|
||||
Parts []struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"parts"`
|
||||
Role string `json:"role"`
|
||||
} `json:"content"`
|
||||
FinishReason string `json:"finishReason"`
|
||||
Index int `json:"index"`
|
||||
SafetyRatings []struct {
|
||||
Category string `json:"category"`
|
||||
Probability string `json:"probability"`
|
||||
} `json:"safetyRatings"`
|
||||
} `json:"candidates"`
|
||||
PromptFeedback struct {
|
||||
SafetyRatings []struct {
|
||||
Category string `json:"category"`
|
||||
Probability string `json:"probability"`
|
||||
} `json:"safetyRatings"`
|
||||
} `json:"promptFeedback"`
|
||||
Error struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Status string `json:"status"`
|
||||
Details []struct {
|
||||
Type string `json:"@type"`
|
||||
FieldViolations []struct {
|
||||
Field string `json:"field"`
|
||||
Description string `json:"description"`
|
||||
} `json:"fieldViolations"`
|
||||
} `json:"details"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
func ChatProxy(c *gin.Context, chatReq *openai.ChatCompletionRequest) {
|
||||
usagelog := store.Tokens{Model: chatReq.Model}
|
||||
|
||||
token, _ := c.Get("localuser")
|
||||
|
||||
lu, err := store.GetUserByToken(token.(string))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": gin.H{
|
||||
"message": err.Error(),
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
usagelog.UserID = int(lu.ID)
|
||||
var prompts []genai.Part
|
||||
var prompt string
|
||||
for _, msg := range chatReq.Messages {
|
||||
switch ct := msg.Content.(type) {
|
||||
case string:
|
||||
prompt += "<" + msg.Role + ">: " + msg.Content.(string) + "\n"
|
||||
prompts = append(prompts, genai.Text("<"+msg.Role+">: "+msg.Content.(string)))
|
||||
case []any:
|
||||
for _, item := range ct {
|
||||
if m, ok := item.(map[string]interface{}); ok {
|
||||
if m["type"] == "text" {
|
||||
prompt += "<" + msg.Role + ">: " + m["text"].(string) + "\n"
|
||||
prompts = append(prompts, genai.Text("<"+msg.Role+">: "+m["text"].(string)))
|
||||
} else if m["type"] == "image_url" {
|
||||
if url, ok := m["image_url"].(map[string]interface{}); ok {
|
||||
if strings.HasPrefix(url["url"].(string), "http") {
|
||||
fmt.Println("网络图片:", url["url"].(string))
|
||||
} else if strings.HasPrefix(url["url"].(string), "data:image") {
|
||||
fmt.Println("base64:", url["url"].(string)[:20])
|
||||
var mime string
|
||||
// openai 会以 data:image 开头，则去掉 data:image/png;base64, 和 data:image/jpeg;base64,
|
||||
if strings.HasPrefix(url["url"].(string), "data:image/png") {
|
||||
mime = "image/png"
|
||||
} else if strings.HasPrefix(url["url"].(string), "data:image/jpeg") {
|
||||
mime = "image/jpeg"
|
||||
} else {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Unsupported image format"})
|
||||
return
|
||||
}
|
||||
imageString := strings.Split(url["url"].(string), ",")[1]
|
||||
imageBytes, err := base64.StdEncoding.DecodeString(imageString)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
prompts = append(prompts, genai.Blob{MIMEType: mime, Data: imageBytes})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
default:
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"error": gin.H{
|
||||
"message": "Invalid content type",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
if len(chatReq.Tools) > 0 {
|
||||
tooljson, _ := json.Marshal(chatReq.Tools)
|
||||
prompt += "<tools>: " + string(tooljson) + "\n"
|
||||
}
|
||||
}
|
||||
|
||||
usagelog.PromptCount = tokenizer.NumTokensFromStr(prompt, chatReq.Model)
|
||||
|
||||
onekey, err := store.SelectKeyCache("google")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
client, err := genai.NewClient(ctx, option.WithAPIKey(onekey.Key))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
model := client.GenerativeModel(chatReq.Model)
|
||||
model.Tools = []*genai.Tool{}
|
||||
|
||||
iter := model.GenerateContentStream(ctx, prompts...)
|
||||
datachan := make(chan string)
|
||||
// closechan := make(chan error)
|
||||
var result string
|
||||
go func() {
|
||||
for {
|
||||
resp, err := iter.Next()
|
||||
if err == iterator.Done {
|
||||
|
||||
var chatResp openai.ChatCompletionStreamResponse
|
||||
chatResp.Model = chatReq.Model
|
||||
choice := openai.Choice{}
|
||||
choice.FinishReason = "stop"
|
||||
chatResp.Choices = append(chatResp.Choices, choice)
|
||||
datachan <- "data: " + string(chatResp.ByteJson())
|
||||
close(datachan)
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
log.Println(err)
|
||||
var errResp openai.ErrResponse
|
||||
errResp.Error.Code = "500"
|
||||
errResp.Error.Message = err.Error()
|
||||
datachan <- string(errResp.ByteJson())
|
||||
close(datachan)
|
||||
break
|
||||
}
|
||||
var content string
|
||||
if resp.Candidates != nil && len(resp.Candidates) > 0 && len(resp.Candidates[0].Content.Parts) > 0 {
|
||||
if s, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
|
||||
content = string(s)
|
||||
result += content
|
||||
}
|
||||
} else {
|
||||
continue
|
||||
}
|
||||
|
||||
var chatResp openai.ChatCompletionStreamResponse
|
||||
chatResp.Model = chatReq.Model
|
||||
choice := openai.Choice{}
|
||||
choice.Delta.Role = resp.Candidates[0].Content.Role
|
||||
choice.Delta.Content = content
|
||||
chatResp.Choices = append(chatResp.Choices, choice)
|
||||
|
||||
chunk := "data: " + string(chatResp.ByteJson()) + "\n\n"
|
||||
datachan <- chunk
|
||||
}
|
||||
}()
|
||||
|
||||
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() {
|
||||
|
||||
}()
|
||||
return false
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user