81 lines
2.0 KiB
Go
81 lines
2.0 KiB
Go
package controller
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"opencatd-open/internal/dto"
|
|
"opencatd-open/internal/model"
|
|
"opencatd-open/llm"
|
|
"opencatd-open/llm/claude/v2"
|
|
"opencatd-open/llm/google/v2"
|
|
"opencatd-open/llm/openai_compatible"
|
|
"opencatd-open/pkg/tokenizer"
|
|
"strconv"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func (h *Proxy) ChatHandler(c *gin.Context) {
|
|
var chatreq llm.ChatRequest
|
|
if err := c.ShouldBindJSON(&chatreq); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
err := h.SelectApiKey(chatreq.Model)
|
|
if err != nil {
|
|
dto.WrapErrorAsOpenAI(c, 500, err.Error())
|
|
return
|
|
}
|
|
|
|
var llm llm.LLM
|
|
switch *h.apikey.ApiType {
|
|
case "claude":
|
|
llm, err = claude.NewClaude(h.apikey)
|
|
case "gemini":
|
|
llm, err = google.NewGemini(c, h.apikey)
|
|
case "openai", "azure", "github":
|
|
fallthrough
|
|
default:
|
|
llm, err = openai_compatible.NewOpenAICompatible(h.apikey)
|
|
if err != nil {
|
|
dto.WrapErrorAsOpenAI(c, 500, fmt.Errorf("create llm client error: %w", err).Error())
|
|
return
|
|
}
|
|
}
|
|
|
|
if !chatreq.Stream {
|
|
resp, err := llm.Chat(c, chatreq)
|
|
if err != nil {
|
|
dto.WrapErrorAsOpenAI(c, 500, err.Error())
|
|
}
|
|
c.JSON(http.StatusOK, resp)
|
|
|
|
} else {
|
|
datachan, err := llm.StreamChat(c, chatreq)
|
|
if err != nil {
|
|
dto.WrapErrorAsOpenAI(c, 500, err.Error())
|
|
}
|
|
for data := range datachan {
|
|
c.SSEvent("", data)
|
|
}
|
|
}
|
|
|
|
llmusage := llm.GetTokenUsage()
|
|
|
|
cost := tokenizer.Cost(llmusage.Model, llmusage.PromptTokens+llmusage.ToolsTokens, llmusage.CompletionTokens)
|
|
userid, _ := strconv.ParseInt(c.GetString("user_id"), 10, 64)
|
|
usage := model.Usage{
|
|
UserID: userid,
|
|
Model: llmusage.Model,
|
|
Stream: chatreq.Stream,
|
|
PromptTokens: llmusage.PromptTokens + llmusage.ToolsTokens,
|
|
CompletionTokens: llmusage.CompletionTokens,
|
|
TotalTokens: llmusage.TotalTokens,
|
|
Cost: fmt.Sprintf("%f", cost),
|
|
}
|
|
h.SendUsage(&usage)
|
|
defer fmt.Println("cost:", cost, "prompt_tokens:", llmusage.PromptTokens, "completion_tokens:", llmusage.CompletionTokens, "total_tokens:", llmusage.TotalTokens)
|
|
|
|
}
|