78 lines
1.8 KiB
Go
78 lines
1.8 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"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func (h *Proxy) ChatHandler(c *gin.Context) {
|
|
user := c.MustGet("user").(*model.User)
|
|
if user == nil {
|
|
dto.WrapErrorAsOpenAI(c, 401, "Unauthorized")
|
|
return
|
|
}
|
|
|
|
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()
|
|
llmusage.User = user
|
|
llmusage.TokenID = c.GetInt64("token_id")
|
|
cost := tokenizer.Cost(llmusage.Model, llmusage.PromptTokens+llmusage.ToolsTokens, llmusage.CompletionTokens)
|
|
|
|
h.SendUsage(llmusage)
|
|
defer fmt.Println("cost:", cost, "prompt_tokens:", llmusage.PromptTokens, "completion_tokens:", llmusage.CompletionTokens, "total_tokens:", llmusage.TotalTokens)
|
|
|
|
}
|