refactor: move backend files to backend/ directory
Reorganize project structure: - backend/cmd/openteam/ — entry point - backend/internal/ — core packages - backend/middleware/ — HTTP middleware - backend/router/ — route setup - backend/wire/ — dependency injection - backend/pkg/ — shared utilities - backend/go.mod, go.sum — Go module files Updated Makefile to work from backend/ directory. Removed old lowercase makefile.
This commit is contained in:
@@ -0,0 +1,565 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
"opencatd-open/internal/pkg/apikey"
|
||||
"opencatd-open/internal/pkg/crypto"
|
||||
"opencatd-open/internal/pkg/jwt"
|
||||
"opencatd-open/internal/auth"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type Handler struct {
|
||||
db *gorm.DB
|
||||
userDAO *dao.UserDAO
|
||||
apiKeyDAO *dao.ApiKeyDAO
|
||||
channelDAO *dao.ChannelDAO
|
||||
modelDAO *dao.ModelDAO
|
||||
usageDAO *dao.UsageDAO
|
||||
dailyDAO *dao.DailyUsageDAO
|
||||
}
|
||||
|
||||
func NewHandler(db *gorm.DB) *Handler {
|
||||
return &Handler{
|
||||
db: db,
|
||||
userDAO: dao.NewUserDAO(db),
|
||||
apiKeyDAO: dao.NewApiKeyDAO(db),
|
||||
channelDAO: dao.NewChannelDAO(db),
|
||||
modelDAO: dao.NewModelDAO(db),
|
||||
usageDAO: dao.NewUsageDAO(db),
|
||||
dailyDAO: dao.NewDailyUsageDAO(db),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Auth ---
|
||||
|
||||
func (h *Handler) Register(c *gin.Context) {
|
||||
var req struct {
|
||||
Username string `json:"username" binding:"required"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
Email string `json:"email" binding:"required,email"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Check if first user (becomes admin)
|
||||
var count int64
|
||||
h.db.Model(&store.User{}).Count(&count)
|
||||
|
||||
role := store.RoleUser
|
||||
if count == 0 {
|
||||
role = store.RoleAdmin
|
||||
}
|
||||
|
||||
hash := crypto.Sha256Hex(req.Password)
|
||||
user := &store.User{
|
||||
Username: req.Username,
|
||||
Email: req.Email,
|
||||
PasswordHash: hash,
|
||||
Role: role,
|
||||
Status: store.UserStatusActive,
|
||||
}
|
||||
|
||||
if err := h.userDAO.Create(user); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "username or email already exists"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "registered"})
|
||||
}
|
||||
|
||||
func (h *Handler) Login(c *gin.Context) {
|
||||
var req struct {
|
||||
Username string `json:"username" binding:"required"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
user, err := h.userDAO.GetByUsername(req.Username)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid credentials"})
|
||||
return
|
||||
}
|
||||
|
||||
hash := crypto.Sha256Hex(req.Password)
|
||||
if user.PasswordHash != hash {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid credentials"})
|
||||
return
|
||||
}
|
||||
|
||||
secret := auth.GetSecretKey()
|
||||
accessToken, refreshToken, err := jwt.GenerateTokenPair(user.ID, user.Username, user.Role, secret, 24*time.Hour, 7*24*time.Hour)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate token"})
|
||||
return
|
||||
}
|
||||
|
||||
// Update last login
|
||||
now := time.Now()
|
||||
user.LastLoginAt = &now
|
||||
h.userDAO.Update(user)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": 200,
|
||||
"data": gin.H{
|
||||
"token": accessToken,
|
||||
"access_token": accessToken,
|
||||
"refresh_token": refreshToken,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (h *Handler) Me(c *gin.Context) {
|
||||
userID, _ := c.Get("user_id")
|
||||
user, err := h.userDAO.GetByID(userID.(uint64))
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "user not found"})
|
||||
return
|
||||
}
|
||||
// Map role string to number for frontend compatibility
|
||||
roleNum := 1 // default user
|
||||
if user.Role == store.RoleAdmin {
|
||||
roleNum = 10
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"code": 200,
|
||||
"data": gin.H{
|
||||
"id": user.ID,
|
||||
"username": user.Username,
|
||||
"email": user.Email,
|
||||
"role": roleNum,
|
||||
"status": user.Status,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// --- Users ---
|
||||
|
||||
func (h *Handler) ListUsers(c *gin.Context) {
|
||||
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "20"))
|
||||
offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0"))
|
||||
users, total, err := h.userDAO.List(limit, offset)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": users, "total": total})
|
||||
}
|
||||
|
||||
func (h *Handler) CreateUser(c *gin.Context) {
|
||||
var req struct {
|
||||
Username string `json:"username" binding:"required"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
Email string `json:"email" binding:"required,email"`
|
||||
Role string `json:"role"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
role := store.RoleUser
|
||||
if req.Role != "" {
|
||||
role = req.Role
|
||||
}
|
||||
|
||||
hash := crypto.Sha256Hex(req.Password)
|
||||
user := &store.User{
|
||||
Username: req.Username,
|
||||
Email: req.Email,
|
||||
PasswordHash: hash,
|
||||
Role: role,
|
||||
Status: store.UserStatusActive,
|
||||
}
|
||||
|
||||
if err := h.userDAO.Create(user); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "username or email already exists"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, user)
|
||||
}
|
||||
|
||||
func (h *Handler) DeleteUser(c *gin.Context) {
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
||||
return
|
||||
}
|
||||
if err := h.userDAO.Delete(id); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "deleted"})
|
||||
}
|
||||
|
||||
// --- API Keys ---
|
||||
|
||||
func (h *Handler) ListApiKeys(c *gin.Context) {
|
||||
userID, _ := c.Get("user_id")
|
||||
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "20"))
|
||||
offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0"))
|
||||
keys, total, err := h.apiKeyDAO.ListByUserID(userID.(uint64), limit, offset)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": keys, "total": total})
|
||||
}
|
||||
|
||||
func (h *Handler) CreateApiKey(c *gin.Context) {
|
||||
var req struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
|
||||
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
userID, _ := c.Get("user_id")
|
||||
keyValue, _ := apikey.Generate()
|
||||
|
||||
key := &store.APIKey{
|
||||
UserID: userID.(uint64),
|
||||
Name: req.Name,
|
||||
KeyHash: apikey.Hash(keyValue),
|
||||
KeyPrefix: keyValue[:8],
|
||||
QuotaTokensPerDay: req.QuotaTokensPerDay,
|
||||
QuotaRequestsPerDay: req.QuotaRequestsPerDay,
|
||||
Status: store.KeyStatusActive,
|
||||
}
|
||||
|
||||
if err := h.apiKeyDAO.Create(key); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"key": keyValue,
|
||||
"id": key.ID,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *Handler) DeleteApiKey(c *gin.Context) {
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
||||
return
|
||||
}
|
||||
if err := h.apiKeyDAO.Delete(id); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "deleted"})
|
||||
}
|
||||
|
||||
// --- Channels ---
|
||||
|
||||
func (h *Handler) ListChannels(c *gin.Context) {
|
||||
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "20"))
|
||||
offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0"))
|
||||
channels, total, err := h.channelDAO.List(limit, offset)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": channels, "total": total})
|
||||
}
|
||||
|
||||
func (h *Handler) CreateChannel(c *gin.Context) {
|
||||
var req struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Provider string `json:"provider" binding:"required"`
|
||||
BaseURL string `json:"base_url" binding:"required"`
|
||||
APIKey string `json:"api_key" binding:"required"`
|
||||
Priority int `json:"priority"`
|
||||
Weight int `json:"weight"`
|
||||
Formats []string `json:"formats"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
encrypted, err := crypto.Encrypt(req.APIKey)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to encrypt API key"})
|
||||
return
|
||||
}
|
||||
|
||||
if req.Weight == 0 {
|
||||
req.Weight = 1
|
||||
}
|
||||
|
||||
ch := &store.Channel{
|
||||
Name: req.Name,
|
||||
Provider: req.Provider,
|
||||
BaseURL: req.BaseURL,
|
||||
APIKeyEnc: encrypted,
|
||||
Weight: req.Weight,
|
||||
Priority: req.Priority,
|
||||
Formats: req.Formats,
|
||||
Enabled: true,
|
||||
}
|
||||
|
||||
if err := h.channelDAO.Create(ch); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "channel name already exists"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, ch)
|
||||
}
|
||||
|
||||
func (h *Handler) UpdateChannel(c *gin.Context) {
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
||||
return
|
||||
}
|
||||
|
||||
ch, err := h.channelDAO.GetByID(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "channel not found"})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"base_url"`
|
||||
APIKey string `json:"api_key"`
|
||||
Priority *int `json:"priority"`
|
||||
Weight *int `json:"weight"`
|
||||
Formats []string `json:"formats"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if req.Name != "" {
|
||||
ch.Name = req.Name
|
||||
}
|
||||
if req.BaseURL != "" {
|
||||
ch.BaseURL = req.BaseURL
|
||||
}
|
||||
if req.APIKey != "" {
|
||||
encrypted, err := crypto.Encrypt(req.APIKey)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to encrypt API key"})
|
||||
return
|
||||
}
|
||||
ch.APIKeyEnc = encrypted
|
||||
}
|
||||
if req.Priority != nil {
|
||||
ch.Priority = *req.Priority
|
||||
}
|
||||
if req.Weight != nil {
|
||||
ch.Weight = *req.Weight
|
||||
}
|
||||
if req.Formats != nil {
|
||||
ch.Formats = req.Formats
|
||||
}
|
||||
if req.Enabled != nil {
|
||||
ch.Enabled = *req.Enabled
|
||||
}
|
||||
|
||||
if err := h.channelDAO.Update(ch); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, ch)
|
||||
}
|
||||
|
||||
func (h *Handler) DeleteChannel(c *gin.Context) {
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
||||
return
|
||||
}
|
||||
if err := h.channelDAO.Delete(id); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "deleted"})
|
||||
}
|
||||
|
||||
// --- Models ---
|
||||
|
||||
func (h *Handler) ListModels(c *gin.Context) {
|
||||
limit, _ := strconv.Atoi(c.DefaultQuery("limit", "20"))
|
||||
offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0"))
|
||||
models, total, err := h.modelDAO.List(limit, offset)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"data": models, "total": total})
|
||||
}
|
||||
|
||||
func (h *Handler) CreateModel(c *gin.Context) {
|
||||
var req struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
DisplayName string `json:"display_name"`
|
||||
InputPrice float64 `json:"input_price"`
|
||||
OutputPrice float64 `json:"output_price"`
|
||||
CacheReadPrice float64 `json:"cache_read_price"`
|
||||
Sort int `json:"sort"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
m := &store.Model{
|
||||
Name: req.Name,
|
||||
DisplayName: req.DisplayName,
|
||||
InputPrice: req.InputPrice,
|
||||
OutputPrice: req.OutputPrice,
|
||||
CacheReadPrice: req.CacheReadPrice,
|
||||
Sort: req.Sort,
|
||||
Enabled: true,
|
||||
}
|
||||
|
||||
if err := h.modelDAO.Create(m); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "model name already exists"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, m)
|
||||
}
|
||||
|
||||
func (h *Handler) UpdateModel(c *gin.Context) {
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
||||
return
|
||||
}
|
||||
|
||||
m, err := h.modelDAO.GetByID(id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "model not found"})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
DisplayName string `json:"display_name"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
Sort *int `json:"sort"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if req.DisplayName != "" {
|
||||
m.DisplayName = req.DisplayName
|
||||
}
|
||||
if req.InputPrice != nil {
|
||||
m.InputPrice = *req.InputPrice
|
||||
}
|
||||
if req.OutputPrice != nil {
|
||||
m.OutputPrice = *req.OutputPrice
|
||||
}
|
||||
if req.CacheReadPrice != nil {
|
||||
m.CacheReadPrice = *req.CacheReadPrice
|
||||
}
|
||||
if req.Sort != nil {
|
||||
m.Sort = *req.Sort
|
||||
}
|
||||
if req.Enabled != nil {
|
||||
m.Enabled = *req.Enabled
|
||||
}
|
||||
|
||||
if err := h.modelDAO.Update(m); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, m)
|
||||
}
|
||||
|
||||
func (h *Handler) DeleteModel(c *gin.Context) {
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
||||
return
|
||||
}
|
||||
if err := h.modelDAO.Delete(id); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"message": "deleted"})
|
||||
}
|
||||
|
||||
// --- Channel-Model Bindings ---
|
||||
|
||||
func (h *Handler) BindChannelModels(c *gin.Context) {
|
||||
channelID, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid channel id"})
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Bindings []struct {
|
||||
ModelID uint64 `json:"model_id"`
|
||||
UpstreamModel string `json:"upstream_model"`
|
||||
Weight int `json:"weight"`
|
||||
} `json:"bindings"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
bindings := make([]store.ChannelModelBinding, len(req.Bindings))
|
||||
for i, b := range req.Bindings {
|
||||
bindings[i] = store.ChannelModelBinding{
|
||||
ChannelID: channelID,
|
||||
ModelID: b.ModelID,
|
||||
UpstreamModel: b.UpstreamModel,
|
||||
Weight: b.Weight,
|
||||
}
|
||||
}
|
||||
|
||||
if err := h.channelDAO.BindModels(channelID, bindings); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "bound"})
|
||||
}
|
||||
|
||||
func (h *Handler) GetChannelModels(c *gin.Context) {
|
||||
channelID, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid channel id"})
|
||||
return
|
||||
}
|
||||
|
||||
bindings, err := h.channelDAO.GetChannelModels(channelID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"data": bindings})
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"opencatd-open/internal/store"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
type Claims struct {
|
||||
UserID uint64 `json:"user_id"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
type TokenPair struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
RefreshToken string `json:"refresh_token,omitempty"`
|
||||
}
|
||||
|
||||
func GenerateTokenPair(user *store.User, secret string, accessExpire, refreshExpire time.Duration) (*TokenPair, error) {
|
||||
accessToken, err := generateToken(user, "access", secret, accessExpire)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
refreshToken, err := generateToken(user, "refresh", secret, refreshExpire)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &TokenPair{
|
||||
AccessToken: accessToken,
|
||||
RefreshToken: refreshToken,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func generateToken(user *store.User, tokenType, secret string, expire time.Duration) (string, error) {
|
||||
now := time.Now()
|
||||
claims := Claims{
|
||||
UserID: user.ID,
|
||||
Name: user.Username,
|
||||
Type: tokenType,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(now.Add(expire)),
|
||||
IssuedAt: jwt.NewNumericDate(now),
|
||||
NotBefore: jwt.NewNumericDate(now),
|
||||
},
|
||||
}
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
return token.SignedString([]byte(secret))
|
||||
}
|
||||
|
||||
func ValidateToken(tokenString, secret string) (*Claims, error) {
|
||||
token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
|
||||
if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, errors.New("unexpected signing method")
|
||||
}
|
||||
return []byte(secret), nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if claims, ok := token.Claims.(*Claims); ok && token.Valid {
|
||||
return claims, nil
|
||||
}
|
||||
return nil, jwt.ErrInvalidKey
|
||||
}
|
||||
|
||||
// GetSecretKey returns the JWT secret key from environment or config
|
||||
func GetSecretKey() string {
|
||||
secret := os.Getenv("SECRET_KEY")
|
||||
if secret == "" {
|
||||
secret = "default-secret-key-change-in-production"
|
||||
}
|
||||
return secret
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
package channel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"math/rand"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
"opencatd-open/internal/pkg/crypto"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
channelDAO *dao.ChannelDAO
|
||||
modelDAO *dao.ModelDAO
|
||||
|
||||
// Health tracking
|
||||
mu sync.RWMutex
|
||||
healthStatus map[uint64]*channelHealth
|
||||
}
|
||||
|
||||
type channelHealth struct {
|
||||
status string
|
||||
consecutive int
|
||||
lastCheck time.Time
|
||||
cooldown time.Time
|
||||
}
|
||||
|
||||
func NewService(channelDAO *dao.ChannelDAO, modelDAO *dao.ModelDAO) *Service {
|
||||
return &Service{
|
||||
channelDAO: channelDAO,
|
||||
modelDAO: modelDAO,
|
||||
healthStatus: make(map[uint64]*channelHealth),
|
||||
}
|
||||
}
|
||||
|
||||
// SelectChannel selects the best channel for a given model using weighted random selection
|
||||
func (s *Service) SelectChannel(ctx context.Context, modelName string) (*store.Channel, error) {
|
||||
channels, err := s.channelDAO.GetEnabledChannelsByModel(modelName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get channels for model %s: %w", modelName, err)
|
||||
}
|
||||
if len(channels) == 0 {
|
||||
return nil, fmt.Errorf("no enabled channels for model: %s", modelName)
|
||||
}
|
||||
|
||||
// Filter out unhealthy channels
|
||||
candidates := s.filterHealthy(channels)
|
||||
if len(candidates) == 0 {
|
||||
// If all channels are unhealthy, try the first one anyway
|
||||
candidates = channels[:1]
|
||||
}
|
||||
|
||||
// Weighted random selection
|
||||
totalWeight := 0
|
||||
for _, ch := range candidates {
|
||||
totalWeight += ch.Weight
|
||||
}
|
||||
if totalWeight == 0 {
|
||||
return candidates[0], nil
|
||||
}
|
||||
|
||||
r := rand.Intn(totalWeight)
|
||||
for _, ch := range candidates {
|
||||
r -= ch.Weight
|
||||
if r < 0 {
|
||||
return ch, nil
|
||||
}
|
||||
}
|
||||
|
||||
return candidates[0], nil
|
||||
}
|
||||
|
||||
// GetChannelByKeyID decrypts the API key for a channel
|
||||
func (s *Service) GetChannelByKeyID(ctx context.Context, channelID uint64) (*store.Channel, error) {
|
||||
ch, err := s.channelDAO.GetByID(channelID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// GetAPIKey decrypts the channel's API key
|
||||
func (s *Service) GetAPIKey(ch *store.Channel) (string, error) {
|
||||
return crypto.Decrypt(ch.APIKeyEnc)
|
||||
}
|
||||
|
||||
// RecordSuccess records a successful request to a channel
|
||||
func (s *Service) RecordSuccess(channelID uint64) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
h := s.getOrCreateHealth(channelID)
|
||||
h.consecutive = 0
|
||||
h.status = store.ChannelHealthHealthy
|
||||
h.lastCheck = time.Now()
|
||||
}
|
||||
|
||||
// RecordFailure records a failed request to a channel
|
||||
func (s *Service) RecordFailure(channelID uint64) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
h := s.getOrCreateHealth(channelID)
|
||||
h.consecutive++
|
||||
h.lastCheck = time.Now()
|
||||
|
||||
if h.consecutive >= 3 {
|
||||
h.status = store.ChannelHealthDegraded
|
||||
h.cooldown = time.Now().Add(5 * time.Minute)
|
||||
}
|
||||
if h.consecutive >= 5 {
|
||||
h.status = store.ChannelHealthCooldown
|
||||
h.cooldown = time.Now().Add(15 * time.Minute)
|
||||
}
|
||||
}
|
||||
|
||||
// RecordTimeout records a timeout to a channel
|
||||
func (s *Service) RecordTimeout(channelID uint64) {
|
||||
s.RecordFailure(channelID)
|
||||
}
|
||||
|
||||
func (s *Service) getOrCreateHealth(channelID uint64) *channelHealth {
|
||||
h, ok := s.healthStatus[channelID]
|
||||
if !ok {
|
||||
h = &channelHealth{
|
||||
status: store.ChannelHealthHealthy,
|
||||
}
|
||||
s.healthStatus[channelID] = h
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
func (s *Service) filterHealthy(channels []*store.Channel) []*store.Channel {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
var healthy []*store.Channel
|
||||
now := time.Now()
|
||||
|
||||
for _, ch := range channels {
|
||||
h, ok := s.healthStatus[ch.ID]
|
||||
if !ok {
|
||||
healthy = append(healthy, ch)
|
||||
continue
|
||||
}
|
||||
|
||||
// Check if cooldown has expired
|
||||
if now.After(h.cooldown) && h.cooldown.IsZero() == false {
|
||||
h.consecutive = 0
|
||||
h.status = store.ChannelHealthHealthy
|
||||
healthy = append(healthy, ch)
|
||||
continue
|
||||
}
|
||||
|
||||
if h.status == store.ChannelHealthHealthy || h.status == store.ChannelHealthDegraded {
|
||||
healthy = append(healthy, ch)
|
||||
}
|
||||
}
|
||||
|
||||
return healthy
|
||||
}
|
||||
|
||||
// GetHealthStatus returns the health status of a channel
|
||||
func (s *Service) GetHealthStatus(channelID uint64) string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
h, ok := s.healthStatus[channelID]
|
||||
if !ok {
|
||||
return store.ChannelHealthHealthy
|
||||
}
|
||||
return h.status
|
||||
}
|
||||
|
||||
// ChannelCandidate represents a channel with its resolved API key
|
||||
type ChannelCandidate struct {
|
||||
Channel *store.Channel
|
||||
APIKey string
|
||||
Format string
|
||||
}
|
||||
|
||||
// SelectCandidates returns candidates for a model, sorted by priority
|
||||
func (s *Service) SelectCandidates(ctx context.Context, modelName string, preferredFormat string) ([]ChannelCandidate, error) {
|
||||
channels, err := s.channelDAO.GetEnabledChannelsByModel(modelName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var candidates []ChannelCandidate
|
||||
for _, ch := range channels {
|
||||
// Check if channel supports the preferred format
|
||||
formats := ch.FormatsEffective()
|
||||
supported := false
|
||||
for _, f := range formats {
|
||||
if f == preferredFormat || preferredFormat == "" {
|
||||
supported = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !supported {
|
||||
continue
|
||||
}
|
||||
|
||||
apiKey, err := crypto.Decrypt(ch.APIKeyEnc)
|
||||
if err != nil {
|
||||
log.Printf("Failed to decrypt API key for channel %s: %v", ch.Name, err)
|
||||
continue
|
||||
}
|
||||
|
||||
candidates = append(candidates, ChannelCandidate{
|
||||
Channel: ch,
|
||||
APIKey: apiKey,
|
||||
Format: preferredFormat,
|
||||
})
|
||||
}
|
||||
|
||||
return candidates, nil
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package channel
|
||||
|
||||
import (
|
||||
"opencatd-open/internal/store"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestChannelFormatsEffective(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
channel store.Channel
|
||||
expected []string
|
||||
}{
|
||||
{
|
||||
name: "anthropic default",
|
||||
channel: store.Channel{
|
||||
Provider: store.ChannelProviderAnthropic,
|
||||
},
|
||||
expected: []string{store.FormatMessages},
|
||||
},
|
||||
{
|
||||
name: "openai default",
|
||||
channel: store.Channel{
|
||||
Provider: store.ChannelProviderOpenAI,
|
||||
},
|
||||
expected: []string{store.FormatChat, store.FormatResponses},
|
||||
},
|
||||
{
|
||||
name: "compatible default",
|
||||
channel: store.Channel{
|
||||
Provider: store.ChannelProviderCompatible,
|
||||
},
|
||||
expected: []string{store.FormatChat},
|
||||
},
|
||||
{
|
||||
name: "custom formats override",
|
||||
channel: store.Channel{
|
||||
Provider: store.ChannelProviderOpenAI,
|
||||
Formats: []string{store.FormatChat},
|
||||
},
|
||||
expected: []string{store.FormatChat},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tt.channel.FormatsEffective()
|
||||
if len(result) != len(tt.expected) {
|
||||
t.Errorf("FormatsEffective() returned %d formats, want %d", len(result), len(tt.expected))
|
||||
return
|
||||
}
|
||||
for i, f := range result {
|
||||
if f != tt.expected[i] {
|
||||
t.Errorf("FormatsEffective()[%d] = %q, want %q", i, f, tt.expected[i])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChannelUpstreamURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
channel store.Channel
|
||||
proto string
|
||||
path string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "basic openai",
|
||||
channel: store.Channel{
|
||||
BaseURL: "https://api.openai.com",
|
||||
},
|
||||
proto: "chat",
|
||||
path: "/chat/completions",
|
||||
expected: "https://api.openai.com/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "with trailing slash",
|
||||
channel: store.Channel{
|
||||
BaseURL: "https://api.openai.com/",
|
||||
},
|
||||
proto: "chat",
|
||||
path: "/chat/completions",
|
||||
expected: "https://api.openai.com/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "with version segment",
|
||||
channel: store.Channel{
|
||||
BaseURL: "https://api.openai.com/v1",
|
||||
},
|
||||
proto: "chat",
|
||||
path: "/chat/completions",
|
||||
expected: "https://api.openai.com/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "custom base URL per protocol",
|
||||
channel: store.Channel{
|
||||
BaseURL: "https://default.openai.com",
|
||||
BaseURLs: map[string]string{"chat": "https://chat.openai.com"},
|
||||
},
|
||||
proto: "chat",
|
||||
path: "/chat/completions",
|
||||
expected: "https://chat.openai.com/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "empty base",
|
||||
channel: store.Channel{
|
||||
BaseURL: "",
|
||||
},
|
||||
proto: "chat",
|
||||
path: "/chat/completions",
|
||||
expected: "/chat/completions",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tt.channel.UpstreamURL(tt.proto, tt.path)
|
||||
if result != tt.expected {
|
||||
t.Errorf("UpstreamURL() = %q, want %q", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package channel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
"opencatd-open/internal/pkg/crypto"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
type HealthChecker struct {
|
||||
channelDAO *dao.ChannelDAO
|
||||
service *Service
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func NewHealthChecker(channelDAO *dao.ChannelDAO, service *Service) *HealthChecker {
|
||||
return &HealthChecker{
|
||||
channelDAO: channelDAO,
|
||||
service: service,
|
||||
client: &http.Client{
|
||||
Timeout: 10 * time.Second,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// CheckChannel performs a health check on a channel
|
||||
func (hc *HealthChecker) CheckChannel(ctx context.Context, channel *store.Channel) error {
|
||||
apiKey, err := crypto.Decrypt(channel.APIKeyEnc)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to decrypt API key: %w", err)
|
||||
}
|
||||
|
||||
// Simple health check: try to list models
|
||||
var url string
|
||||
switch channel.Provider {
|
||||
case store.ChannelProviderOpenAI:
|
||||
url = channel.UpstreamURL("chat", "/models")
|
||||
case store.ChannelProviderAnthropic:
|
||||
url = "https://api.anthropic.com/v1/models"
|
||||
default:
|
||||
url = channel.UpstreamURL("chat", "/models")
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
// Set headers based on provider
|
||||
switch channel.Provider {
|
||||
case store.ChannelProviderOpenAI, store.ChannelProviderCompatible:
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
case store.ChannelProviderAnthropic:
|
||||
req.Header.Set("x-api-key", apiKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := hc.client.Do(req)
|
||||
if err != nil {
|
||||
hc.service.RecordFailure(channel.ID)
|
||||
return fmt.Errorf("health check failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode == http.StatusOK {
|
||||
hc.service.RecordSuccess(channel.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
hc.service.RecordFailure(channel.ID)
|
||||
return fmt.Errorf("health check returned status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
// CheckAllChannels checks health of all enabled channels
|
||||
func (hc *HealthChecker) CheckAllChannels(ctx context.Context) error {
|
||||
channels, err := hc.channelDAO.ListEnabled()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, ch := range channels {
|
||||
if err := hc.CheckChannel(ctx, ch); err != nil {
|
||||
fmt.Printf("Channel %s health check failed: %v\n", ch.Name, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StartPeriodicCheck starts periodic health checks
|
||||
func (hc *HealthChecker) StartPeriodicCheck(ctx context.Context, interval time.Duration) {
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := hc.CheckAllChannels(ctx); err != nil {
|
||||
fmt.Printf("Periodic health check error: %v\n", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"opencatd-open/internal/store"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
var LoadCmd = &cobra.Command{
|
||||
Use: "load",
|
||||
Short: "import user.json -> db",
|
||||
Long: "\nimport user.json -> db",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
db := store.DB
|
||||
var cont int64
|
||||
if err := db.Model(&store.User{}).Count(&cont).Error; err != nil {
|
||||
fmt.Println(err)
|
||||
return
|
||||
}
|
||||
if cont == 0 {
|
||||
fmt.Println("创建管理员之后再操作")
|
||||
return
|
||||
}
|
||||
if _, err := os.Stat("./db/user.json"); os.IsNotExist(err) {
|
||||
log.Fatalln("404! user.json is not found.")
|
||||
return
|
||||
}
|
||||
file, err := os.Open("./db/user.json")
|
||||
if err != nil {
|
||||
fmt.Println("Error opening file:", err)
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
var usermap []map[string]string
|
||||
if err := json.NewDecoder(file).Decode(&usermap); err != nil {
|
||||
fmt.Println("解析文件失败:", err)
|
||||
return
|
||||
}
|
||||
for _, um := range usermap {
|
||||
name := um["username"]
|
||||
if name == "" {
|
||||
name = um["name"]
|
||||
}
|
||||
if name == "" {
|
||||
fmt.Println("获取不到数据")
|
||||
continue
|
||||
}
|
||||
_ = "sk-ot-" + strings.ReplaceAll(uuid.New().String(), "-", "")
|
||||
fmt.Printf("Import user: %s\n", name)
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
var SaveCmd = &cobra.Command{
|
||||
Use: "save",
|
||||
Short: "backup user info -> user.json",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"opencatd-open/internal/service"
|
||||
"opencatd-open/pkg/config"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type Api struct {
|
||||
cfg *config.Config
|
||||
db *gorm.DB
|
||||
userService *service.UserServiceImpl
|
||||
tokenService *service.TokenServiceImpl
|
||||
keyService *service.ApiKeyServiceImpl
|
||||
webAuthService *service.WebAuthnService
|
||||
usageService *service.UsageService
|
||||
}
|
||||
|
||||
func NewApi(cfg *config.Config, db *gorm.DB, userService *service.UserServiceImpl, tokenService *service.TokenServiceImpl, keyService *service.ApiKeyServiceImpl, webAuthService *service.WebAuthnService, usageService *service.UsageService) *Api {
|
||||
return &Api{
|
||||
cfg: cfg,
|
||||
db: db,
|
||||
userService: userService,
|
||||
tokenService: tokenService,
|
||||
keyService: keyService,
|
||||
webAuthService: webAuthService,
|
||||
usageService: usageService,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"opencatd-open/internal/channel"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
"opencatd-open/pkg/config"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type Proxy struct {
|
||||
ctx context.Context
|
||||
cfg *config.Config
|
||||
db *gorm.DB
|
||||
wg *sync.WaitGroup
|
||||
httpClient *http.Client
|
||||
|
||||
userDAO *dao.UserDAO
|
||||
apiKeyDAO *dao.ApiKeyDAO
|
||||
usageDAO *dao.UsageDAO
|
||||
dailyDAO *dao.DailyUsageDAO
|
||||
channelSvc *channel.Service
|
||||
}
|
||||
|
||||
func NewProxy(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.WaitGroup, userDAO *dao.UserDAO, apiKeyDAO *dao.ApiKeyDAO, usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Proxy {
|
||||
client := http.DefaultClient
|
||||
if os.Getenv("LOCAL_PROXY") != "" {
|
||||
proxyUrl, err := url.Parse(os.Getenv("LOCAL_PROXY"))
|
||||
if err == nil {
|
||||
tr := &http.Transport{
|
||||
Proxy: http.ProxyURL(proxyUrl),
|
||||
}
|
||||
client.Transport = tr
|
||||
}
|
||||
}
|
||||
|
||||
np := &Proxy{
|
||||
ctx: ctx,
|
||||
cfg: cfg,
|
||||
db: db,
|
||||
wg: wg,
|
||||
httpClient: client,
|
||||
userDAO: userDAO,
|
||||
apiKeyDAO: apiKeyDAO,
|
||||
usageDAO: usageDAO,
|
||||
dailyDAO: dailyDAO,
|
||||
}
|
||||
|
||||
return np
|
||||
}
|
||||
|
||||
// SetChannelService sets the channel service (called after construction)
|
||||
func (p *Proxy) SetChannelService(svc *channel.Service) {
|
||||
p.channelSvc = svc
|
||||
}
|
||||
|
||||
func (p *Proxy) HandleProxy(c *gin.Context) {
|
||||
path := c.Request.URL.Path
|
||||
switch {
|
||||
case path == "/v1/chat/completions":
|
||||
// TODO: Phase 3 - implement chat completions handler
|
||||
c.JSON(http.StatusNotImplemented, gin.H{"error": "chat completions not yet implemented"})
|
||||
case strings.HasPrefix(path, "/v1/messages"):
|
||||
// TODO: Phase 3 - implement messages handler
|
||||
c.JSON(http.StatusNotImplemented, gin.H{"error": "messages not yet implemented"})
|
||||
case path == "/v1/responses":
|
||||
// TODO: Phase 3 - implement responses handler
|
||||
c.JSON(http.StatusNotImplemented, gin.H{"error": "responses not yet implemented"})
|
||||
default:
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "unknown endpoint"})
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Proxy) HandleModels(c *gin.Context) {
|
||||
// TODO: Phase 3 - implement models list
|
||||
c.JSON(http.StatusOK, gin.H{"object": "list", "data": []interface{}{}})
|
||||
}
|
||||
|
||||
func (p *Proxy) GetDB() *gorm.DB {
|
||||
return p.db
|
||||
}
|
||||
|
||||
// SelectChannel selects the best channel for a model
|
||||
func (p *Proxy) SelectChannel(modelName string) (*store.Channel, error) {
|
||||
if p.channelSvc == nil {
|
||||
return nil, fmt.Errorf("channel service not initialized")
|
||||
}
|
||||
return p.channelSvc.SelectChannel(p.ctx, modelName)
|
||||
}
|
||||
|
||||
// RecordSuccess records a successful request
|
||||
func (p *Proxy) RecordSuccess(channelID uint64) {
|
||||
if p.channelSvc != nil {
|
||||
p.channelSvc.RecordSuccess(channelID)
|
||||
}
|
||||
}
|
||||
|
||||
// RecordFailure records a failed request
|
||||
func (p *Proxy) RecordFailure(channelID uint64) {
|
||||
if p.channelSvc != nil {
|
||||
p.channelSvc.RecordFailure(channelID)
|
||||
}
|
||||
}
|
||||
|
||||
// SendUsagePlaceholder placeholder for usage processing
|
||||
func (p *Proxy) SendUsagePlaceholder(model string, userID uint64, promptTokens, completionTokens int) {
|
||||
log.Printf("Usage: model=%s user=%d prompt=%d completion=%d", model, userID, promptTokens, completionTokens)
|
||||
}
|
||||
|
||||
// Placeholder to keep the file compilable
|
||||
var _ = json.Marshal
|
||||
var _ = io.ReadAll
|
||||
@@ -0,0 +1,68 @@
|
||||
package dao
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"opencatd-open/internal/store"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type ApiKeyDAO struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewApiKeyDAO(db *gorm.DB) *ApiKeyDAO {
|
||||
return &ApiKeyDAO{db: db}
|
||||
}
|
||||
|
||||
func (d *ApiKeyDAO) Create(apiKey *store.APIKey) error {
|
||||
if apiKey == nil {
|
||||
return errors.New("apiKey is nil")
|
||||
}
|
||||
return d.db.Create(apiKey).Error
|
||||
}
|
||||
|
||||
func (d *ApiKeyDAO) GetByID(id uint64) (*store.APIKey, error) {
|
||||
var apiKey store.APIKey
|
||||
err := d.db.First(&apiKey, id).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &apiKey, nil
|
||||
}
|
||||
|
||||
func (d *ApiKeyDAO) GetByHash(keyHash string) (*store.APIKey, error) {
|
||||
var apiKey store.APIKey
|
||||
err := d.db.Where("key_hash = ? AND status = ?", keyHash, store.KeyStatusActive).First(&apiKey).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &apiKey, nil
|
||||
}
|
||||
|
||||
func (d *ApiKeyDAO) ListByUserID(userID uint64, limit, offset int) ([]*store.APIKey, int64, error) {
|
||||
var apiKeys []*store.APIKey
|
||||
var total int64
|
||||
db := d.db.Where("user_id = ?", userID)
|
||||
db.Model(&store.APIKey{}).Count(&total)
|
||||
err := db.Limit(limit).Offset(offset).Order("created_at DESC").Find(&apiKeys).Error
|
||||
return apiKeys, total, err
|
||||
}
|
||||
|
||||
func (d *ApiKeyDAO) Update(apiKey *store.APIKey) error {
|
||||
if apiKey == nil {
|
||||
return errors.New("apiKey is nil")
|
||||
}
|
||||
return d.db.Save(apiKey).Error
|
||||
}
|
||||
|
||||
func (d *ApiKeyDAO) Delete(id uint64) error {
|
||||
return d.db.Delete(&store.APIKey{}, id).Error
|
||||
}
|
||||
|
||||
func (d *ApiKeyDAO) BatchDelete(ids []uint64) error {
|
||||
if len(ids) == 0 {
|
||||
return errors.New("ids is empty")
|
||||
}
|
||||
return d.db.Delete(&store.APIKey{}, ids).Error
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package dao
|
||||
|
||||
import (
|
||||
"opencatd-open/internal/store"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type ChannelDAO struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewChannelDAO(db *gorm.DB) *ChannelDAO {
|
||||
return &ChannelDAO{db: db}
|
||||
}
|
||||
|
||||
func (d *ChannelDAO) Create(channel *store.Channel) error {
|
||||
return d.db.Create(channel).Error
|
||||
}
|
||||
|
||||
func (d *ChannelDAO) GetByID(id uint64) (*store.Channel, error) {
|
||||
var channel store.Channel
|
||||
err := d.db.First(&channel, id).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &channel, nil
|
||||
}
|
||||
|
||||
func (d *ChannelDAO) GetByName(name string) (*store.Channel, error) {
|
||||
var channel store.Channel
|
||||
err := d.db.Where("name = ?", name).First(&channel).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &channel, nil
|
||||
}
|
||||
|
||||
func (d *ChannelDAO) List(limit, offset int) ([]*store.Channel, int64, error) {
|
||||
var channels []*store.Channel
|
||||
var total int64
|
||||
d.db.Model(&store.Channel{}).Count(&total)
|
||||
err := d.db.Limit(limit).Offset(offset).Order("priority DESC, weight DESC").Find(&channels).Error
|
||||
return channels, total, err
|
||||
}
|
||||
|
||||
func (d *ChannelDAO) ListEnabled() ([]*store.Channel, error) {
|
||||
var channels []*store.Channel
|
||||
err := d.db.Where("enabled = ?", true).Order("priority DESC, weight DESC").Find(&channels).Error
|
||||
return channels, err
|
||||
}
|
||||
|
||||
func (d *ChannelDAO) Update(channel *store.Channel) error {
|
||||
return d.db.Save(channel).Error
|
||||
}
|
||||
|
||||
func (d *ChannelDAO) Delete(id uint64) error {
|
||||
return d.db.Delete(&store.Channel{}, id).Error
|
||||
}
|
||||
|
||||
// BindModels binds models to a channel (replaces existing bindings)
|
||||
func (d *ChannelDAO) BindModels(channelID uint64, bindings []store.ChannelModelBinding) error {
|
||||
return d.db.Transaction(func(tx *gorm.DB) error {
|
||||
// Delete existing bindings
|
||||
if err := tx.Where("channel_id = ?", channelID).Delete(&store.ChannelModelBinding{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
// Create new bindings
|
||||
for i := range bindings {
|
||||
bindings[i].ChannelID = channelID
|
||||
}
|
||||
return tx.Create(&bindings).Error
|
||||
})
|
||||
}
|
||||
|
||||
// GetChannelModels returns all models bound to a channel
|
||||
func (d *ChannelDAO) GetChannelModels(channelID uint64) ([]store.ChannelModelBinding, error) {
|
||||
var bindings []store.ChannelModelBinding
|
||||
err := d.db.Where("channel_id = ?", channelID).Find(&bindings).Error
|
||||
return bindings, err
|
||||
}
|
||||
|
||||
// GetModelChannels returns all channels that support a given model (by model name)
|
||||
func (d *ChannelDAO) GetModelChannels(modelName string) ([]store.ChannelModelBinding, error) {
|
||||
var bindings []store.ChannelModelBinding
|
||||
err := d.db.
|
||||
Joins("JOIN channels ON channels.id = channel_model_bindings.channel_id").
|
||||
Joins("JOIN models ON models.id = channel_model_bindings.model_id").
|
||||
Where("models.name = ? AND channels.enabled = ?", modelName, true).
|
||||
Find(&bindings).Error
|
||||
return bindings, err
|
||||
}
|
||||
|
||||
// GetEnabledChannelsByModel returns enabled channels for a model, ordered by priority/weight
|
||||
func (d *ChannelDAO) GetEnabledChannelsByModel(modelName string) ([]*store.Channel, error) {
|
||||
var channels []*store.Channel
|
||||
err := d.db.
|
||||
Distinct("channels.*").
|
||||
Joins("JOIN channel_model_bindings ON channel_model_bindings.channel_id = channels.id").
|
||||
Joins("JOIN models ON models.id = channel_model_bindings.model_id").
|
||||
Where("models.name = ? AND channels.enabled = ?", modelName, true).
|
||||
Order("channels.priority DESC, channels.weight DESC").
|
||||
Find(&channels).Error
|
||||
return channels, err
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package dao
|
||||
|
||||
import (
|
||||
"opencatd-open/internal/store"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type ModelDAO struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewModelDAO(db *gorm.DB) *ModelDAO {
|
||||
return &ModelDAO{db: db}
|
||||
}
|
||||
|
||||
func (d *ModelDAO) Create(model *store.Model) error {
|
||||
return d.db.Create(model).Error
|
||||
}
|
||||
|
||||
func (d *ModelDAO) GetByID(id uint64) (*store.Model, error) {
|
||||
var model store.Model
|
||||
err := d.db.First(&model, id).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &model, nil
|
||||
}
|
||||
|
||||
func (d *ModelDAO) GetByName(name string) (*store.Model, error) {
|
||||
var model store.Model
|
||||
err := d.db.Where("name = ?", name).First(&model).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &model, nil
|
||||
}
|
||||
|
||||
func (d *ModelDAO) List(limit, offset int) ([]*store.Model, int64, error) {
|
||||
var models []*store.Model
|
||||
var total int64
|
||||
d.db.Model(&store.Model{}).Count(&total)
|
||||
err := d.db.Limit(limit).Offset(offset).Order("sort ASC, name ASC").Find(&models).Error
|
||||
return models, total, err
|
||||
}
|
||||
|
||||
func (d *ModelDAO) ListEnabled() ([]*store.Model, error) {
|
||||
var models []*store.Model
|
||||
err := d.db.Where("enabled = ?", true).Order("sort ASC, name ASC").Find(&models).Error
|
||||
return models, err
|
||||
}
|
||||
|
||||
func (d *ModelDAO) Update(model *store.Model) error {
|
||||
return d.db.Save(model).Error
|
||||
}
|
||||
|
||||
func (d *ModelDAO) Delete(id uint64) error {
|
||||
return d.db.Delete(&store.Model{}, id).Error
|
||||
}
|
||||
|
||||
// Upsert creates or updates a model by name
|
||||
func (d *ModelDAO) Upsert(model *store.Model) error {
|
||||
return d.db.Where("name = ?", model.Name).Assign(store.Model{
|
||||
DisplayName: model.DisplayName,
|
||||
InputPrice: model.InputPrice,
|
||||
OutputPrice: model.OutputPrice,
|
||||
CacheReadPrice: model.CacheReadPrice,
|
||||
Enabled: model.Enabled,
|
||||
Sort: model.Sort,
|
||||
}).FirstOrCreate(model).Error
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package dao
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"opencatd-open/internal/store"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type TokenDAO struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewTokenDAO(db *gorm.DB) *TokenDAO {
|
||||
return &TokenDAO{db: db}
|
||||
}
|
||||
|
||||
func (d *TokenDAO) GetByKey(key string) (*store.User, error) {
|
||||
var user store.User
|
||||
err := d.db.Where("username = ?", key).First(&user).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (d *TokenDAO) GetByID(id uint64) (*store.User, error) {
|
||||
var user store.User
|
||||
err := d.db.First(&user, id).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
// Placeholder to avoid compile errors - will be expanded in Phase 1
|
||||
var _ = errors.New
|
||||
var _ = gorm.ErrRecordNotFound
|
||||
@@ -0,0 +1,99 @@
|
||||
package dao
|
||||
|
||||
import (
|
||||
"context"
|
||||
"opencatd-open/internal/store"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type UsageDAO struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
type DailyUsageDAO struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewUsageDAO(db *gorm.DB) *UsageDAO {
|
||||
return &UsageDAO{db: db}
|
||||
}
|
||||
|
||||
func NewDailyUsageDAO(db *gorm.DB) *DailyUsageDAO {
|
||||
return &DailyUsageDAO{db: db}
|
||||
}
|
||||
|
||||
// UsageLog DAO
|
||||
func (d *UsageDAO) Create(ctx context.Context, log *store.UsageLog) error {
|
||||
return d.db.WithContext(ctx).Create(log).Error
|
||||
}
|
||||
|
||||
func (d *UsageDAO) BatchCreate(ctx context.Context, logs []*store.UsageLog) error {
|
||||
return d.db.WithContext(ctx).Create(logs).Error
|
||||
}
|
||||
|
||||
func (d *UsageDAO) ListByUserID(ctx context.Context, userID uint64, limit, offset int) ([]*store.UsageLog, error) {
|
||||
var logs []*store.UsageLog
|
||||
err := d.db.WithContext(ctx).
|
||||
Where("user_id = ?", userID).
|
||||
Order("created_at DESC").
|
||||
Limit(limit).
|
||||
Offset(offset).
|
||||
Find(&logs).Error
|
||||
return logs, err
|
||||
}
|
||||
|
||||
func (d *UsageDAO) Delete(ctx context.Context, id uint64) error {
|
||||
return d.db.WithContext(ctx).Delete(&store.UsageLog{}, id).Error
|
||||
}
|
||||
|
||||
func (d *UsageDAO) CountByUserID(ctx context.Context, userID uint64) (int64, error) {
|
||||
var count int64
|
||||
err := d.db.WithContext(ctx).Model(&store.UsageLog{}).Where("user_id = ?", userID).Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// UsageDaily DAO
|
||||
func (d *DailyUsageDAO) Create(ctx context.Context, log *store.UsageDaily) error {
|
||||
return d.db.WithContext(ctx).Create(log).Error
|
||||
}
|
||||
|
||||
func (d *DailyUsageDAO) ListByUserID(ctx context.Context, userID uint64, limit, offset int) ([]*store.UsageDaily, error) {
|
||||
var logs []*store.UsageDaily
|
||||
err := d.db.WithContext(ctx).
|
||||
Where("user_id = ?", userID).
|
||||
Order("date DESC").
|
||||
Limit(limit).
|
||||
Offset(offset).
|
||||
Find(&logs).Error
|
||||
return logs, err
|
||||
}
|
||||
|
||||
func (d *DailyUsageDAO) GetByDate(ctx context.Context, userID uint64, date string) (*store.UsageDaily, error) {
|
||||
var log store.UsageDaily
|
||||
err := d.db.WithContext(ctx).
|
||||
Where("user_id = ? AND date = ?", userID, date).
|
||||
First(&log).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &log, nil
|
||||
}
|
||||
|
||||
func (d *DailyUsageDAO) UpsertDailyUsage(ctx context.Context, log *store.UsageDaily) error {
|
||||
return d.db.WithContext(ctx).Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "user_id"}, {Name: "model_id"}, {Name: "date"}},
|
||||
DoUpdates: clause.AssignmentColumns([]string{"requests", "input_tokens", "output_tokens", "cache_read_tokens", "cost"}),
|
||||
}).Create(log).Error
|
||||
}
|
||||
|
||||
func (d *DailyUsageDAO) ListByDateRange(ctx context.Context, userID uint64, start, end time.Time) ([]*store.UsageDaily, error) {
|
||||
var logs []*store.UsageDaily
|
||||
err := d.db.WithContext(ctx).
|
||||
Where("user_id = ? AND date >= ? AND date <= ?", userID, start.Format("2006-01-02"), end.Format("2006-01-02")).
|
||||
Order("date DESC").
|
||||
Find(&logs).Error
|
||||
return logs, err
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package dao
|
||||
|
||||
import (
|
||||
"opencatd-open/internal/store"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type UserDAO struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewUserDAO(db *gorm.DB) *UserDAO {
|
||||
return &UserDAO{db: db}
|
||||
}
|
||||
|
||||
func (d *UserDAO) Create(user *store.User) error {
|
||||
return d.db.Create(user).Error
|
||||
}
|
||||
|
||||
func (d *UserDAO) GetByID(id uint64) (*store.User, error) {
|
||||
var user store.User
|
||||
err := d.db.First(&user, id).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (d *UserDAO) GetByUsername(username string) (*store.User, error) {
|
||||
var user store.User
|
||||
err := d.db.Where("username = ?", username).First(&user).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (d *UserDAO) GetByEmail(email string) (*store.User, error) {
|
||||
var user store.User
|
||||
err := d.db.Where("email = ?", email).First(&user).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &user, nil
|
||||
}
|
||||
|
||||
func (d *UserDAO) List(limit, offset int) ([]*store.User, int64, error) {
|
||||
var users []*store.User
|
||||
var total int64
|
||||
d.db.Model(&store.User{}).Count(&total)
|
||||
err := d.db.Limit(limit).Offset(offset).Order("created_at DESC").Find(&users).Error
|
||||
return users, total, err
|
||||
}
|
||||
|
||||
func (d *UserDAO) Update(user *store.User) error {
|
||||
return d.db.Save(user).Error
|
||||
}
|
||||
|
||||
func (d *UserDAO) Delete(id uint64) error {
|
||||
return d.db.Delete(&store.User{}, id).Error
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
package dto
|
||||
|
||||
type BatchIDRequest struct {
|
||||
UserID *int64 `json:"user_id"`
|
||||
IDs []int64 `json:"ids" binding:"required"`
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
package dto
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Error struct {
|
||||
Code int `json:"code,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
}
|
||||
|
||||
func WrapErrorAsOpenAI(c *gin.Context, code int, msg string) {
|
||||
c.JSON(code, gin.H{
|
||||
"error": Error{
|
||||
Code: code,
|
||||
Message: msg,
|
||||
},
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package dto
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"regexp"
|
||||
"time"
|
||||
|
||||
validation "github.com/go-ozzo/ozzo-validation/v4"
|
||||
)
|
||||
|
||||
// TeamKey 结构体定义
|
||||
type TeamKey struct {
|
||||
ID *int64 `json:"id,omitempty"`
|
||||
UserID *int64 `json:"userID,omitempty"`
|
||||
Name *string `json:"name,omitempty"` // 必须
|
||||
Key *string `json:"key,omitempty"`
|
||||
Status *int64 `json:"status,omitempty"` // 默认1 允许,0禁止
|
||||
Quota *int64 `json:"quota,omitempty"` // UnlimitedQuota不为1 的时候必须
|
||||
UnlimitedQuota *bool `json:"unlimitedQuota,omitempty"` // 默认1 不限制,0限制
|
||||
UsedQuota *int64 `json:"usedQuota,omitempty"`
|
||||
CreatedAt *int64 `json:"createdAt,omitempty"`
|
||||
ExpiredAt *int64 `json:"expiredAt,omitempty"` // 可选
|
||||
}
|
||||
|
||||
// DefaultTeamKey 创建一个具有默认值的 TeamKey
|
||||
func DefaultTeamKey() TeamKey {
|
||||
status := int64(1) // 默认允许
|
||||
unlimitedQuota := true // 默认不限制
|
||||
createdAt := time.Now().Unix()
|
||||
|
||||
return TeamKey{
|
||||
Status: &status,
|
||||
UnlimitedQuota: &unlimitedQuota,
|
||||
CreatedAt: &createdAt,
|
||||
}
|
||||
}
|
||||
|
||||
// Validate 验证 TeamKey 结构体
|
||||
func (t TeamKey) Validate() error {
|
||||
// 自定义验证规则
|
||||
var quotaRule validation.Rule = validation.Skip
|
||||
if t.UnlimitedQuota != nil && !*t.UnlimitedQuota {
|
||||
quotaRule = validation.Required.Error("当 UnlimitedQuota 为 false 时,Quota 是必填项")
|
||||
}
|
||||
|
||||
// 过期时间校验
|
||||
var expiredAtRule validation.Rule = validation.Skip
|
||||
if t.ExpiredAt != nil {
|
||||
expiredAtRule = validation.Min(time.Now().Unix()).Error("过期时间不能早于当前时间")
|
||||
}
|
||||
|
||||
return validation.ValidateStruct(&t,
|
||||
// ID 通常由系统生成,不需要验证
|
||||
|
||||
// UserID 可选,但如果提供必须大于 0
|
||||
validation.Field(&t.UserID,
|
||||
validation.When(t.UserID != nil, validation.Min(int64(1)).Error("用户 ID 必须大于 0"))),
|
||||
|
||||
// Name 是必填字段
|
||||
validation.Field(&t.Name,
|
||||
validation.Required.Error("名称不能为空"),
|
||||
validation.When(t.Name != nil, validation.Length(1, 100).Error("名称长度应在 1-100 之间"))),
|
||||
|
||||
// Key 可选,但如果提供需要符合特定格式
|
||||
validation.Field(&t.Key,
|
||||
validation.When(t.Key != nil,
|
||||
validation.Length(1, 255).Error("Key 长度应在 1-255 之间")),
|
||||
validation.Match(regexp.MustCompile(`^[^\s]+$`)).Error("Key 不能包含空格"),
|
||||
),
|
||||
|
||||
// Status 只能是 0 或 1
|
||||
validation.Field(&t.Status,
|
||||
validation.When(t.Status != nil, validation.In(int64(0), int64(1)).Error("状态只能是 0(禁止) 或 1(允许)"))),
|
||||
|
||||
// Quota 要求依赖于 UnlimitedQuota
|
||||
validation.Field(&t.Quota, quotaRule,
|
||||
validation.When(t.Quota != nil, validation.Min(int64(1)).Error("配额必须大于 0"))),
|
||||
|
||||
// UnlimitedQuota 是否限制配额
|
||||
validation.Field(&t.UnlimitedQuota),
|
||||
|
||||
// UsedQuota 系统维护,不需要验证
|
||||
validation.Field(&t.UsedQuota,
|
||||
validation.When(t.UsedQuota != nil, validation.Min(int64(0)).Error("已使用配额不能为负数"))),
|
||||
|
||||
// CreatedAt 系统维护,不需要验证
|
||||
validation.Field(&t.CreatedAt),
|
||||
|
||||
// ExpiredAt 可选,但如果提供必须大于当前时间
|
||||
validation.Field(&t.ExpiredAt, expiredAtRule),
|
||||
)
|
||||
}
|
||||
|
||||
// ValidateCreate 创建时的特殊验证
|
||||
func (t TeamKey) ValidateCreate() error {
|
||||
// 首先进行基本验证
|
||||
if err := t.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 创建时的额外验证
|
||||
if t.Name == nil {
|
||||
return errors.New("创建时必须提供名称")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
package dto
|
||||
|
||||
type Passkey struct {
|
||||
ID int64 `json:"id" gorm:"column:id;primaryKey;autoIncrement"`
|
||||
Name string `json:"name" gorm:"column:name"` // 凭证名称,用于用户识别不同的设备
|
||||
SignCount uint32 `json:"sign_count" gorm:"column:sign_count"` // 签名计数器,用于防止重放攻击
|
||||
DeviceType string `json:"device_type" gorm:"column:device_type"` // 设备类型,如"platform"或"cross-platform"
|
||||
LastUsedAt int64 `json:"last_used_at" gorm:"column:last_used_at"` // 最后使用时间
|
||||
CreatedAt int64 `json:"created_at,omitempty" gorm:"autoCreateTime"`
|
||||
UpdatedAt int64 `json:"updated_at,omitempty" gorm:"autoUpdateTime"`
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package dto
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Result struct {
|
||||
Code int `json:"code"`
|
||||
Msg string `json:"msg"`
|
||||
Data any `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
func Success(ctx *gin.Context, data any) {
|
||||
ctx.JSON(http.StatusOK, Result{
|
||||
Code: 200,
|
||||
Data: data,
|
||||
Msg: "success",
|
||||
})
|
||||
}
|
||||
|
||||
func Fail(c *gin.Context, code int, err string) {
|
||||
c.AbortWithStatusJSON(code, gin.H{
|
||||
"code": code,
|
||||
"error": err,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package dto
|
||||
|
||||
type User struct {
|
||||
Username string `json:"username" binding:"required,min=3,max=32"`
|
||||
Password string `json:"password" binding:"required,min=4"`
|
||||
}
|
||||
|
||||
type Auth struct {
|
||||
Token string `json:"token"`
|
||||
ExpiresIn int64 `json:"expires_in"`
|
||||
}
|
||||
|
||||
type ChangePassword struct {
|
||||
Password string `json:"password" binding:"required,min=4"`
|
||||
NewPassword string `json:"newpassword" binding:"required,min=4"`
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"opencatd-open/internal/pkg/crypto"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const Prefix = "sk-ot-"
|
||||
|
||||
// Generate 生成新的 API Key,返回明文和哈希
|
||||
func Generate() (plaintext, hash string) {
|
||||
b := make([]byte, 24)
|
||||
_, _ = rand.Read(b)
|
||||
raw := hex.EncodeToString(b)
|
||||
plaintext = Prefix + raw
|
||||
hash = crypto.Sha256Hex(plaintext)
|
||||
return
|
||||
}
|
||||
|
||||
// Valid 校验 API Key 格式
|
||||
func Valid(key string) bool {
|
||||
return strings.HasPrefix(key, Prefix)
|
||||
}
|
||||
|
||||
// Hash 计算 API Key 的 SHA-256 哈希
|
||||
func Hash(key string) string {
|
||||
return crypto.Sha256Hex(key)
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package crypto
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
)
|
||||
|
||||
func defaultKey() []byte {
|
||||
key := os.Getenv("ENCRYPT_KEY")
|
||||
if key == "" {
|
||||
key = "opencatd-default-key-change-me"
|
||||
}
|
||||
h := sha256.Sum256([]byte(key))
|
||||
return h[:] // 32 bytes
|
||||
}
|
||||
|
||||
// Encrypt encrypts plaintext using AES-GCM with the default key
|
||||
func Encrypt(plaintext string) (string, error) {
|
||||
enc, err := NewEncryptor(defaultKey())
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return enc.Encrypt(plaintext)
|
||||
}
|
||||
|
||||
// Decrypt decrypts ciphertext using AES-GCM with the default key
|
||||
func Decrypt(encoded string) (string, error) {
|
||||
enc, err := NewEncryptor(defaultKey())
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return enc.Decrypt(encoded)
|
||||
}
|
||||
|
||||
// Sha256Hex is a convenience wrapper for SHA-256 hex hashing
|
||||
func Sha256Hex(data string) string {
|
||||
h := sha256.Sum256([]byte(data))
|
||||
return hex.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
// Encryptor AES-GCM 加密器
|
||||
type Encryptor struct {
|
||||
key []byte
|
||||
}
|
||||
|
||||
// NewEncryptor 创建加密器(key 为 16/24/32 字节)
|
||||
func NewEncryptor(key []byte) (*Encryptor, error) {
|
||||
switch len(key) {
|
||||
case 16, 24, 32:
|
||||
default:
|
||||
return nil, errors.New("crypto: invalid key length, must be 16, 24, or 32 bytes")
|
||||
}
|
||||
return &Encryptor{key: key}, nil
|
||||
}
|
||||
|
||||
// Encrypt AES-GCM 加密,返回 base64 编码的密文
|
||||
func (e *Encryptor) Encrypt(plaintext string) (string, error) {
|
||||
block, err := aes.NewCipher(e.key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
nonce := make([]byte, gcm.NonceSize())
|
||||
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
|
||||
return "", err
|
||||
}
|
||||
ciphertext := gcm.Seal(nonce, nonce, []byte(plaintext), nil)
|
||||
return base64.StdEncoding.EncodeToString(ciphertext), nil
|
||||
}
|
||||
|
||||
// Decrypt AES-GCM 解密
|
||||
func (e *Encryptor) Decrypt(encoded string) (string, error) {
|
||||
data, err := base64.StdEncoding.DecodeString(encoded)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
block, err := aes.NewCipher(e.key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
nonceSize := gcm.NonceSize()
|
||||
if len(data) < nonceSize {
|
||||
return "", errors.New("ciphertext too short")
|
||||
}
|
||||
nonce, ciphertext := data[:nonceSize], data[nonceSize:]
|
||||
plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(plaintext), nil
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package crypto
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestEncryptDecrypt(t *testing.T) {
|
||||
plaintext := "sk-test-api-key-12345"
|
||||
|
||||
encrypted, err := Encrypt(plaintext)
|
||||
if err != nil {
|
||||
t.Fatalf("Encrypt() error = %v", err)
|
||||
}
|
||||
|
||||
if encrypted == plaintext {
|
||||
t.Error("Encrypt() returned plaintext")
|
||||
}
|
||||
|
||||
decrypted, err := Decrypt(encrypted)
|
||||
if err != nil {
|
||||
t.Fatalf("Decrypt() error = %v", err)
|
||||
}
|
||||
|
||||
if decrypted != plaintext {
|
||||
t.Errorf("Decrypt() = %q, want %q", decrypted, plaintext)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSha256Hex(t *testing.T) {
|
||||
input := "test"
|
||||
result := Sha256Hex(input)
|
||||
|
||||
if len(result) != 64 {
|
||||
t.Errorf("Sha256Hex() returned %d chars, want 64", len(result))
|
||||
}
|
||||
|
||||
// Same input should produce same hash
|
||||
result2 := Sha256Hex(input)
|
||||
if result != result2 {
|
||||
t.Error("Sha256Hex() not deterministic")
|
||||
}
|
||||
|
||||
// Different input should produce different hash
|
||||
result3 := Sha256Hex("different")
|
||||
if result == result3 {
|
||||
t.Error("Sha256Hex() same hash for different inputs")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncryptorInvalidKey(t *testing.T) {
|
||||
_, err := NewEncryptor([]byte("short"))
|
||||
if err == nil {
|
||||
t.Error("NewEncryptor() should error with invalid key length")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package jwt
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
gojwt "github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
type Claims struct {
|
||||
UserID uint64 `json:"user_id"`
|
||||
Name string `json:"name"`
|
||||
Role string `json:"role"`
|
||||
gojwt.RegisteredClaims
|
||||
}
|
||||
|
||||
// GenerateTokenPair 生成 access + refresh token
|
||||
func GenerateTokenPair(userID uint64, name, role, secret string, accessExpire, refreshExpire time.Duration) (accessToken, refreshToken string, err error) {
|
||||
accessToken, err = generateToken(userID, name, role, "access", secret, accessExpire)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
refreshToken, err = generateToken(userID, name, role, "refresh", secret, refreshExpire)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func generateToken(userID uint64, name, role, tokenType, secret string, expire time.Duration) (string, error) {
|
||||
now := time.Now()
|
||||
claims := Claims{
|
||||
UserID: userID,
|
||||
Name: name,
|
||||
Role: role,
|
||||
RegisteredClaims: gojwt.RegisteredClaims{
|
||||
ExpiresAt: gojwt.NewNumericDate(now.Add(expire)),
|
||||
IssuedAt: gojwt.NewNumericDate(now),
|
||||
NotBefore: gojwt.NewNumericDate(now),
|
||||
},
|
||||
}
|
||||
token := gojwt.NewWithClaims(gojwt.SigningMethodHS256, claims)
|
||||
return token.SignedString([]byte(secret))
|
||||
}
|
||||
|
||||
// ValidateToken 校验 JWT
|
||||
func ValidateToken(tokenString, secret string) (*Claims, error) {
|
||||
token, err := gojwt.ParseWithClaims(tokenString, &Claims{}, func(token *gojwt.Token) (interface{}, error) {
|
||||
if _, ok := token.Method.(*gojwt.SigningMethodHMAC); !ok {
|
||||
return nil, errors.New("unexpected signing method")
|
||||
}
|
||||
return []byte(secret), nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if claims, ok := token.Claims.(*Claims); ok && token.Valid {
|
||||
return claims, nil
|
||||
}
|
||||
return nil, gojwt.ErrInvalidKey
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package ratelimit
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Limiter 内存限流器
|
||||
type Limiter struct {
|
||||
mu sync.Mutex
|
||||
|
||||
// 每用户每秒请求数
|
||||
userRPS map[uint64]*tokenBucket
|
||||
|
||||
// 密钥每日请求计数
|
||||
keyDailyReq map[uint64]*dailyCounter
|
||||
|
||||
// 密钥每日 token 计数
|
||||
keyDailyTokens map[uint64]*dailyCounter
|
||||
}
|
||||
|
||||
type tokenBucket struct {
|
||||
tokens float64
|
||||
maxTokens float64
|
||||
refillRate float64
|
||||
lastRefill time.Time
|
||||
}
|
||||
|
||||
type dailyCounter struct {
|
||||
date string
|
||||
count int64
|
||||
}
|
||||
|
||||
func New() *Limiter {
|
||||
return &Limiter{
|
||||
userRPS: make(map[uint64]*tokenBucket),
|
||||
keyDailyReq: make(map[uint64]*dailyCounter),
|
||||
keyDailyTokens: make(map[uint64]*dailyCounter),
|
||||
}
|
||||
}
|
||||
|
||||
// AllowRequest 检查用户级每秒请求限制
|
||||
func (l *Limiter) AllowRequest(userID uint64, rps int) bool {
|
||||
if rps <= 0 {
|
||||
return true
|
||||
}
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
bucket, ok := l.userRPS[userID]
|
||||
if !ok {
|
||||
bucket = &tokenBucket{
|
||||
tokens: float64(rps),
|
||||
maxTokens: float64(rps),
|
||||
refillRate: float64(rps),
|
||||
lastRefill: time.Now(),
|
||||
}
|
||||
l.userRPS[userID] = bucket
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
elapsed := now.Sub(bucket.lastRefill).Seconds()
|
||||
bucket.tokens += elapsed * bucket.refillRate
|
||||
if bucket.tokens > bucket.maxTokens {
|
||||
bucket.tokens = bucket.maxTokens
|
||||
}
|
||||
bucket.lastRefill = now
|
||||
|
||||
if bucket.tokens < 1 {
|
||||
return false
|
||||
}
|
||||
bucket.tokens--
|
||||
return true
|
||||
}
|
||||
|
||||
// AllowRequestDaily 检查密钥每日请求配额
|
||||
func (l *Limiter) AllowRequestDaily(keyID uint64, quota int) bool {
|
||||
if quota <= 0 {
|
||||
return true
|
||||
}
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
today := time.Now().UTC().Format("2006-01-02")
|
||||
counter, ok := l.keyDailyReq[keyID]
|
||||
if !ok || counter.date != today {
|
||||
l.keyDailyReq[keyID] = &dailyCounter{date: today, count: 1}
|
||||
return true
|
||||
}
|
||||
if counter.count >= int64(quota) {
|
||||
return false
|
||||
}
|
||||
counter.count++
|
||||
return true
|
||||
}
|
||||
|
||||
// TokensUsed 返回密钥今日 token 用量
|
||||
func (l *Limiter) TokensUsed(keyID uint64) int64 {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
today := time.Now().UTC().Format("2006-01-02")
|
||||
counter, ok := l.keyDailyTokens[keyID]
|
||||
if !ok || counter.date != today {
|
||||
return 0
|
||||
}
|
||||
return counter.count
|
||||
}
|
||||
|
||||
// AddTokens 累加密钥今日 token 用量
|
||||
func (l *Limiter) AddTokens(keyID uint64, tokens int64) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
today := time.Now().UTC().Format("2006-01-02")
|
||||
counter, ok := l.keyDailyTokens[keyID]
|
||||
if !ok || counter.date != today {
|
||||
l.keyDailyTokens[keyID] = &dailyCounter{date: today, count: tokens}
|
||||
return
|
||||
}
|
||||
counter.count += tokens
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package resp
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// Error 按 OpenAI 格式返回错误
|
||||
func Error(c *gin.Context, status int, message string) {
|
||||
c.AbortWithStatusJSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"message": message,
|
||||
"type": "api_error",
|
||||
"param": nil,
|
||||
"code": nil,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// ErrorWithType 按 OpenAI 格式返回带类型的错误
|
||||
func ErrorWithType(c *gin.Context, status int, errType, code, message string) {
|
||||
c.AbortWithStatusJSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"message": message,
|
||||
"type": errType,
|
||||
"param": nil,
|
||||
"code": code,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// ErrorAsAnthropic 按 Anthropic 格式返回错误
|
||||
func ErrorAsAnthropic(c *gin.Context, status int, errType, message string) {
|
||||
c.AbortWithStatusJSON(status, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{
|
||||
"type": errType,
|
||||
"message": message,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// OK 返回成功 JSON
|
||||
func OK(c *gin.Context, data interface{}) {
|
||||
c.JSON(http.StatusOK, data)
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/pkoukk/tiktoken-go"
|
||||
)
|
||||
|
||||
// Count 计算字符串的 token 数量
|
||||
func Count(text, model string) int {
|
||||
tkm, err := tiktoken.EncodingForModel(model)
|
||||
if err != nil {
|
||||
tkm, _ = tiktoken.GetEncoding("cl100k_base")
|
||||
}
|
||||
return len(tkm.Encode(text, nil, nil))
|
||||
}
|
||||
|
||||
// Cost 计算模型调用成本(USD,按每百万 token 定价)
|
||||
func Cost(model string, inputTokens, outputTokens int) float64 {
|
||||
var inputPrice, outputPrice float64
|
||||
|
||||
switch {
|
||||
case strings.Contains(model, "gpt-4o-mini"):
|
||||
inputPrice = 0.15
|
||||
outputPrice = 0.60
|
||||
case strings.Contains(model, "gpt-4o"):
|
||||
inputPrice = 2.50
|
||||
outputPrice = 10.00
|
||||
case strings.Contains(model, "gpt-4-turbo"):
|
||||
inputPrice = 10.00
|
||||
outputPrice = 30.00
|
||||
case strings.Contains(model, "gpt-4"):
|
||||
inputPrice = 30.00
|
||||
outputPrice = 60.00
|
||||
case strings.Contains(model, "gpt-3.5-turbo"):
|
||||
inputPrice = 0.50
|
||||
outputPrice = 1.50
|
||||
case strings.Contains(model, "claude-3-5-sonnet"):
|
||||
inputPrice = 3.00
|
||||
outputPrice = 15.00
|
||||
case strings.Contains(model, "claude-3-opus"):
|
||||
inputPrice = 15.00
|
||||
outputPrice = 75.00
|
||||
case strings.Contains(model, "claude-3-haiku"):
|
||||
inputPrice = 0.25
|
||||
outputPrice = 1.25
|
||||
case strings.Contains(model, "claude"):
|
||||
inputPrice = 8.00
|
||||
outputPrice = 24.00
|
||||
case strings.Contains(model, "gemini-1.5-pro"):
|
||||
inputPrice = 3.50
|
||||
outputPrice = 10.50
|
||||
case strings.Contains(model, "gemini-1.5-flash"):
|
||||
inputPrice = 0.35
|
||||
outputPrice = 0.53
|
||||
case strings.Contains(model, "gemini"):
|
||||
inputPrice = 0.50
|
||||
outputPrice = 1.50
|
||||
default:
|
||||
inputPrice = 0.15
|
||||
outputPrice = 0.60
|
||||
}
|
||||
|
||||
cost := float64(inputTokens)/1e6*inputPrice + float64(outputTokens)/1e6*outputPrice
|
||||
if cost < 0.000001 {
|
||||
cost = 0.000001
|
||||
}
|
||||
return cost
|
||||
}
|
||||
|
||||
// CostWithModel 从数据库模型记录获取定价
|
||||
func CostWithModel(inputTokens, outputTokens, cacheReadTokens, cacheCreationTokens int64, inputPrice, outputPrice, cacheReadPrice float64) float64 {
|
||||
cost := float64(inputTokens)/1e6*inputPrice +
|
||||
float64(outputTokens)/1e6*outputPrice +
|
||||
float64(cacheReadTokens)/1e6*cacheReadPrice +
|
||||
float64(cacheCreationTokens)/1e6*inputPrice*1.25
|
||||
if cost < 0.000001 {
|
||||
cost = 0.000001
|
||||
}
|
||||
return cost
|
||||
}
|
||||
|
||||
func init() {
|
||||
_ = fmt.Sprintf // ensure fmt is used
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package convert
|
||||
|
||||
// ChatCompletionRequest represents an OpenAI Chat Completions request
|
||||
type ChatCompletionRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []Message `json:"messages"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty"`
|
||||
N *int `json:"n,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Stop interface{} `json:"stop,omitempty"`
|
||||
MaxTokens *int `json:"max_tokens,omitempty"`
|
||||
PresencePenalty *float64 `json:"presence_penalty,omitempty"`
|
||||
FrequencyPenalty *float64 `json:"frequency_penalty,omitempty"`
|
||||
LogitBias map[string]int `json:"logit_bias,omitempty"`
|
||||
User string `json:"user,omitempty"`
|
||||
Tools []Tool `json:"tools,omitempty"`
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
ResponseFormat interface{} `json:"response_format,omitempty"`
|
||||
Seed *int `json:"seed,omitempty"`
|
||||
}
|
||||
|
||||
// ChatCompletionResponse represents an OpenAI Chat Completions response
|
||||
type ChatCompletionResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []Choice `json:"choices"`
|
||||
Usage *Usage `json:"usage,omitempty"`
|
||||
SystemFingerprint string `json:"system_fingerprint,omitempty"`
|
||||
}
|
||||
|
||||
type Choice struct {
|
||||
Index int `json:"index"`
|
||||
Message Message `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
// ChatCompletionStreamChunk represents a streaming chunk
|
||||
type ChatCompletionStreamChunk struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []StreamChoice `json:"choices"`
|
||||
Usage *Usage `json:"usage,omitempty"`
|
||||
SystemFingerprint string `json:"system_fingerprint,omitempty"`
|
||||
}
|
||||
|
||||
type StreamChoice struct {
|
||||
Index int `json:"index"`
|
||||
Delta StreamDelta `json:"delta"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
type StreamDelta struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,212 @@
|
||||
package convert
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// ChatToMessages converts a Chat Completions request to Anthropic Messages format
|
||||
func ChatToMessages(req *ChatCompletionRequest) (*MessagesRequest, error) {
|
||||
msgs := make([]Message, 0, len(req.Messages))
|
||||
var systemParts []ContentPart
|
||||
|
||||
for _, m := range req.Messages {
|
||||
if m.Role == "system" {
|
||||
// Extract system message content
|
||||
switch v := m.Content.(type) {
|
||||
case string:
|
||||
systemParts = append(systemParts, ContentPart{
|
||||
Type: "text",
|
||||
Text: v,
|
||||
})
|
||||
case []interface{}:
|
||||
for _, part := range v {
|
||||
if p, ok := part.(map[string]interface{}); ok {
|
||||
if t, ok := p["type"].(string); ok && t == "text" {
|
||||
if text, ok := p["text"].(string); ok {
|
||||
systemParts = append(systemParts, ContentPart{
|
||||
Type: "text",
|
||||
Text: text,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
msgs = append(msgs, m)
|
||||
}
|
||||
|
||||
out := &MessagesRequest{
|
||||
Model: req.Model,
|
||||
Messages: msgs,
|
||||
Stream: req.Stream,
|
||||
}
|
||||
|
||||
if len(systemParts) > 0 {
|
||||
out.System = systemParts
|
||||
}
|
||||
|
||||
if req.MaxTokens != nil {
|
||||
out.MaxTokens = *req.MaxTokens
|
||||
} else {
|
||||
defaultMax := 4096
|
||||
out.MaxTokens = defaultMax
|
||||
}
|
||||
|
||||
if req.Temperature != nil {
|
||||
out.Temperature = req.Temperature
|
||||
}
|
||||
if req.TopP != nil {
|
||||
out.TopP = req.TopP
|
||||
}
|
||||
if req.Tools != nil {
|
||||
out.Tools = req.Tools
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// MessagesToChat converts an Anthropic Messages response to Chat Completions format
|
||||
func MessagesToChat(resp *MessagesResponse) (*ChatCompletionResponse, error) {
|
||||
choices := make([]Choice, 0)
|
||||
|
||||
for _, block := range resp.Content {
|
||||
switch block.Type {
|
||||
case "text":
|
||||
choices = append(choices, Choice{
|
||||
Index: len(choices),
|
||||
Message: Message{
|
||||
Role: "assistant",
|
||||
Content: block.Text,
|
||||
},
|
||||
FinishReason: mapStopReason(resp.StopReason),
|
||||
})
|
||||
case "tool_use":
|
||||
toolCall := ToolCall{
|
||||
ID: block.ID,
|
||||
Type: "function",
|
||||
Function: FunctionCall{
|
||||
Name: block.Name,
|
||||
Arguments: toJSON(block.Input),
|
||||
},
|
||||
}
|
||||
if len(choices) == 0 {
|
||||
choices = append(choices, Choice{
|
||||
Index: 0,
|
||||
Message: Message{
|
||||
Role: "assistant",
|
||||
ToolCalls: []ToolCall{toolCall},
|
||||
},
|
||||
FinishReason: "tool_calls",
|
||||
})
|
||||
} else {
|
||||
choices[0].Message.ToolCalls = append(choices[0].Message.ToolCalls, toolCall)
|
||||
choices[0].FinishReason = "tool_calls"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(choices) == 0 {
|
||||
choices = append(choices, Choice{
|
||||
Index: 0,
|
||||
Message: Message{
|
||||
Role: "assistant",
|
||||
Content: "",
|
||||
},
|
||||
FinishReason: "stop",
|
||||
})
|
||||
}
|
||||
|
||||
return &ChatCompletionResponse{
|
||||
ID: resp.ID,
|
||||
Object: "chat.completion",
|
||||
Model: resp.Model,
|
||||
Choices: choices,
|
||||
Usage: &Usage{
|
||||
PromptTokens: resp.Usage.PromptTokens,
|
||||
CompletionTokens: resp.Usage.CompletionTokens,
|
||||
TotalTokens: resp.Usage.PromptTokens + resp.Usage.CompletionTokens,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// MessagesStreamToChatStream converts Anthropic streaming chunks to Chat Completions format
|
||||
func MessagesStreamToChatStream(anthropicEvents []AnthropicStreamEvent, model string) []ChatCompletionStreamChunk {
|
||||
var chunks []ChatCompletionStreamChunk
|
||||
id := fmt.Sprintf("chatcmpl-%d", len(anthropicEvents))
|
||||
|
||||
for _, event := range anthropicEvents {
|
||||
switch event.Type {
|
||||
case "message_start":
|
||||
// Initial chunk with role
|
||||
chunks = append(chunks, ChatCompletionStreamChunk{
|
||||
ID: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Model: model,
|
||||
Choices: []StreamChoice{{
|
||||
Index: 0,
|
||||
Delta: StreamDelta{
|
||||
Role: "assistant",
|
||||
},
|
||||
}},
|
||||
})
|
||||
case "content_block_delta":
|
||||
if event.Delta != nil && event.Delta.Text != "" {
|
||||
chunks = append(chunks, ChatCompletionStreamChunk{
|
||||
ID: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Model: model,
|
||||
Choices: []StreamChoice{{
|
||||
Index: 0,
|
||||
Delta: StreamDelta{
|
||||
Content: event.Delta.Text,
|
||||
},
|
||||
}},
|
||||
})
|
||||
}
|
||||
case "message_delta":
|
||||
finishReason := "stop"
|
||||
if event.Delta != nil && event.Delta.StopReason != "" {
|
||||
finishReason = mapStopReason(event.Delta.StopReason)
|
||||
}
|
||||
chunk := ChatCompletionStreamChunk{
|
||||
ID: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Model: model,
|
||||
Choices: []StreamChoice{{
|
||||
Index: 0,
|
||||
FinishReason: &finishReason,
|
||||
}},
|
||||
}
|
||||
if event.Usage != nil {
|
||||
chunk.Usage = event.Usage
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
func mapStopReason(reason string) string {
|
||||
switch reason {
|
||||
case "end_turn", "stop_sequence":
|
||||
return "stop"
|
||||
case "tool_use":
|
||||
return "tool_calls"
|
||||
case "max_tokens":
|
||||
return "length"
|
||||
default:
|
||||
return "stop"
|
||||
}
|
||||
}
|
||||
|
||||
func toJSON(v interface{}) string {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return "{}"
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
@@ -0,0 +1,298 @@
|
||||
package convert
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// ChatToResponses converts a Chat Completions request to Responses API format
|
||||
func ChatToResponses(req *ChatCompletionRequest) (*ResponsesRequest, error) {
|
||||
var inputItems []InputItem
|
||||
var instructions string
|
||||
|
||||
for _, m := range req.Messages {
|
||||
if m.Role == "system" {
|
||||
if s, ok := m.Content.(string); ok {
|
||||
if instructions != "" {
|
||||
instructions += "\n\n"
|
||||
}
|
||||
instructions += s
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
item := InputItem{
|
||||
Role: m.Role,
|
||||
Content: m.Content,
|
||||
}
|
||||
inputItems = append(inputItems, item)
|
||||
}
|
||||
|
||||
out := &ResponsesRequest{
|
||||
Model: req.Model,
|
||||
Input: inputItems,
|
||||
Instructions: instructions,
|
||||
Stream: req.Stream,
|
||||
}
|
||||
|
||||
if req.MaxTokens != nil {
|
||||
out.MaxOutputTokens = req.MaxTokens
|
||||
}
|
||||
if req.Temperature != nil {
|
||||
out.Temperature = req.Temperature
|
||||
}
|
||||
if req.TopP != nil {
|
||||
out.TopP = req.TopP
|
||||
}
|
||||
if req.Tools != nil {
|
||||
out.Tools = req.Tools
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ResponsesToChat converts a Responses API response to Chat Completions format
|
||||
func ResponsesToChat(resp *ResponsesResponse) (*ChatCompletionResponse, error) {
|
||||
choices := make([]Choice, 0)
|
||||
|
||||
for _, output := range resp.Output {
|
||||
switch output.Type {
|
||||
case "message":
|
||||
for _, content := range output.Content {
|
||||
switch content.Type {
|
||||
case "output_text":
|
||||
choices = append(choices, Choice{
|
||||
Index: len(choices),
|
||||
Message: Message{
|
||||
Role: "assistant",
|
||||
Content: content.Text,
|
||||
},
|
||||
FinishReason: "stop",
|
||||
})
|
||||
case "function_call":
|
||||
toolCall := ToolCall{
|
||||
ID: content.ID,
|
||||
Type: "function",
|
||||
Function: FunctionCall{
|
||||
Name: content.Name,
|
||||
Arguments: toJSON(content.Input),
|
||||
},
|
||||
}
|
||||
if len(choices) == 0 {
|
||||
choices = append(choices, Choice{
|
||||
Index: 0,
|
||||
Message: Message{
|
||||
Role: "assistant",
|
||||
ToolCalls: []ToolCall{toolCall},
|
||||
},
|
||||
FinishReason: "tool_calls",
|
||||
})
|
||||
} else {
|
||||
choices[0].Message.ToolCalls = append(choices[0].Message.ToolCalls, toolCall)
|
||||
choices[0].FinishReason = "tool_calls"
|
||||
}
|
||||
}
|
||||
}
|
||||
case "function_call_output":
|
||||
// This would be in a user message context
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
if len(choices) == 0 {
|
||||
choices = append(choices, Choice{
|
||||
Index: 0,
|
||||
Message: Message{
|
||||
Role: "assistant",
|
||||
Content: "",
|
||||
},
|
||||
FinishReason: "stop",
|
||||
})
|
||||
}
|
||||
|
||||
return &ChatCompletionResponse{
|
||||
ID: resp.ID,
|
||||
Object: "chat.completion",
|
||||
Model: resp.Model,
|
||||
Choices: choices,
|
||||
Usage: &Usage{
|
||||
PromptTokens: resp.Usage.PromptTokens,
|
||||
CompletionTokens: resp.Usage.CompletionTokens,
|
||||
TotalTokens: resp.Usage.PromptTokens + resp.Usage.CompletionTokens,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ResponsesStreamToChatStream converts Responses API streaming to Chat Completions format
|
||||
func ResponsesStreamToChatStream(events []ResponsesStreamEvent, model string) []ChatCompletionStreamChunk {
|
||||
var chunks []ChatCompletionStreamChunk
|
||||
id := fmt.Sprintf("chatcmpl-%d", len(events))
|
||||
|
||||
for _, event := range events {
|
||||
switch event.Type {
|
||||
case "response.created":
|
||||
chunks = append(chunks, ChatCompletionStreamChunk{
|
||||
ID: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Model: model,
|
||||
Choices: []StreamChoice{{
|
||||
Index: 0,
|
||||
Delta: StreamDelta{
|
||||
Role: "assistant",
|
||||
},
|
||||
}},
|
||||
})
|
||||
case "response.output_item.added":
|
||||
if event.Item != nil && event.Item.Type == "message" {
|
||||
chunks = append(chunks, ChatCompletionStreamChunk{
|
||||
ID: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Model: model,
|
||||
Choices: []StreamChoice{{
|
||||
Index: 0,
|
||||
Delta: StreamDelta{
|
||||
Role: "assistant",
|
||||
},
|
||||
}},
|
||||
})
|
||||
}
|
||||
case "response.content_part.delta":
|
||||
if event.Delta != "" {
|
||||
chunks = append(chunks, ChatCompletionStreamChunk{
|
||||
ID: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Model: model,
|
||||
Choices: []StreamChoice{{
|
||||
Index: 0,
|
||||
Delta: StreamDelta{
|
||||
Content: event.Delta,
|
||||
},
|
||||
}},
|
||||
})
|
||||
}
|
||||
case "response.completed":
|
||||
finishReason := "stop"
|
||||
chunk := ChatCompletionStreamChunk{
|
||||
ID: id,
|
||||
Object: "chat.completion.chunk",
|
||||
Model: model,
|
||||
Choices: []StreamChoice{{
|
||||
Index: 0,
|
||||
FinishReason: &finishReason,
|
||||
}},
|
||||
}
|
||||
chunks = append(chunks, chunk)
|
||||
}
|
||||
}
|
||||
|
||||
return chunks
|
||||
}
|
||||
|
||||
// MessagesToResponses converts an Anthropic Messages request to Responses API format
|
||||
func MessagesToResponses(req *MessagesRequest) (*ResponsesRequest, error) {
|
||||
var inputItems []InputItem
|
||||
var instructions string
|
||||
|
||||
// Handle system message
|
||||
if req.System != nil {
|
||||
switch v := req.System.(type) {
|
||||
case string:
|
||||
instructions = v
|
||||
case []ContentPart:
|
||||
for _, p := range v {
|
||||
if p.Type == "text" {
|
||||
if instructions != "" {
|
||||
instructions += "\n\n"
|
||||
}
|
||||
instructions += p.Text
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, m := range req.Messages {
|
||||
item := InputItem{
|
||||
Role: m.Role,
|
||||
Content: m.Content,
|
||||
}
|
||||
inputItems = append(inputItems, item)
|
||||
}
|
||||
|
||||
out := &ResponsesRequest{
|
||||
Model: req.Model,
|
||||
Input: inputItems,
|
||||
Instructions: instructions,
|
||||
Stream: req.Stream,
|
||||
}
|
||||
|
||||
out.MaxOutputTokens = &req.MaxTokens
|
||||
|
||||
if req.Temperature != nil {
|
||||
out.Temperature = req.Temperature
|
||||
}
|
||||
if req.TopP != nil {
|
||||
out.TopP = req.TopP
|
||||
}
|
||||
if req.Tools != nil {
|
||||
out.Tools = req.Tools
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ResponsesToMessages converts a Responses API response to Anthropic Messages format
|
||||
func ResponsesToMessages(resp *ResponsesResponse) (*MessagesResponse, error) {
|
||||
var content []ContentBlock
|
||||
|
||||
for _, output := range resp.Output {
|
||||
switch output.Type {
|
||||
case "message":
|
||||
for _, c := range output.Content {
|
||||
switch c.Type {
|
||||
case "output_text":
|
||||
content = append(content, ContentBlock{
|
||||
Type: "text",
|
||||
Text: c.Text,
|
||||
})
|
||||
case "function_call":
|
||||
content = append(content, ContentBlock{
|
||||
Type: "tool_use",
|
||||
ID: c.ID,
|
||||
Name: c.Name,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var stopReason string
|
||||
if len(content) > 0 {
|
||||
last := content[len(content)-1]
|
||||
if last.Type == "tool_use" {
|
||||
stopReason = "tool_use"
|
||||
} else {
|
||||
stopReason = "end_turn"
|
||||
}
|
||||
} else {
|
||||
stopReason = "end_turn"
|
||||
}
|
||||
|
||||
return &MessagesResponse{
|
||||
ID: resp.ID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: content,
|
||||
Model: resp.Model,
|
||||
StopReason: stopReason,
|
||||
Usage: resp.Usage,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// toJSON is a helper to convert a value to JSON string
|
||||
func toJSONStr(v interface{}) string {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return "{}"
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
package convert
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestChatToMessages(t *testing.T) {
|
||||
maxTokens := 1024
|
||||
temp := 0.7
|
||||
|
||||
req := &ChatCompletionRequest{
|
||||
Model: "claude-3-sonnet-20240229",
|
||||
Messages: []Message{
|
||||
{Role: "system", Content: "You are a helpful assistant."},
|
||||
{Role: "user", Content: "Hello!"},
|
||||
},
|
||||
MaxTokens: &maxTokens,
|
||||
Temperature: &temp,
|
||||
}
|
||||
|
||||
result, err := ChatToMessages(req)
|
||||
if err != nil {
|
||||
t.Fatalf("ChatToMessages() error = %v", err)
|
||||
}
|
||||
|
||||
if result.Model != "claude-3-sonnet-20240229" {
|
||||
t.Errorf("Model = %q, want %q", result.Model, "claude-3-sonnet-20240229")
|
||||
}
|
||||
|
||||
if len(result.Messages) != 1 {
|
||||
t.Errorf("Messages length = %d, want 1", len(result.Messages))
|
||||
}
|
||||
|
||||
if result.Messages[0].Role != "user" {
|
||||
t.Errorf("Messages[0].Role = %q, want %q", result.Messages[0].Role, "user")
|
||||
}
|
||||
|
||||
if result.System == nil {
|
||||
t.Error("System is nil, want non-nil")
|
||||
}
|
||||
|
||||
if result.MaxTokens != 1024 {
|
||||
t.Errorf("MaxTokens = %d, want 1024", result.MaxTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessagesToChat(t *testing.T) {
|
||||
resp := &MessagesResponse{
|
||||
ID: "msg-123",
|
||||
Model: "claude-3-sonnet-20240229",
|
||||
Content: []ContentBlock{
|
||||
{Type: "text", Text: "Hello! How can I help?"},
|
||||
},
|
||||
StopReason: "end_turn",
|
||||
Usage: Usage{
|
||||
PromptTokens: 10,
|
||||
CompletionTokens: 20,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := MessagesToChat(resp)
|
||||
if err != nil {
|
||||
t.Fatalf("MessagesToChat() error = %v", err)
|
||||
}
|
||||
|
||||
if result.ID != "msg-123" {
|
||||
t.Errorf("ID = %q, want %q", result.ID, "msg-123")
|
||||
}
|
||||
|
||||
if result.Object != "chat.completion" {
|
||||
t.Errorf("Object = %q, want %q", result.Object, "chat.completion")
|
||||
}
|
||||
|
||||
if len(result.Choices) != 1 {
|
||||
t.Errorf("Choices length = %d, want 1", len(result.Choices))
|
||||
return
|
||||
}
|
||||
|
||||
if result.Choices[0].Message.Role != "assistant" {
|
||||
t.Errorf("Choices[0].Message.Role = %q, want %q", result.Choices[0].Message.Role, "assistant")
|
||||
}
|
||||
|
||||
if result.Choices[0].Message.Content != "Hello! How can I help?" {
|
||||
t.Errorf("Choices[0].Message.Content = %q, want %q", result.Choices[0].Message.Content, "Hello! How can I help?")
|
||||
}
|
||||
|
||||
if result.Choices[0].FinishReason != "stop" {
|
||||
t.Errorf("FinishReason = %q, want %q", result.Choices[0].FinishReason, "stop")
|
||||
}
|
||||
|
||||
if result.Usage.TotalTokens != 30 {
|
||||
t.Errorf("Usage.TotalTokens = %d, want 30", result.Usage.TotalTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatToResponses(t *testing.T) {
|
||||
maxTokens := 2048
|
||||
|
||||
req := &ChatCompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Messages: []Message{
|
||||
{Role: "system", Content: "You are a helpful assistant."},
|
||||
{Role: "user", Content: "What is 2+2?"},
|
||||
},
|
||||
MaxTokens: &maxTokens,
|
||||
}
|
||||
|
||||
result, err := ChatToResponses(req)
|
||||
if err != nil {
|
||||
t.Fatalf("ChatToResponses() error = %v", err)
|
||||
}
|
||||
|
||||
if result.Model != "gpt-4o" {
|
||||
t.Errorf("Model = %q, want %q", result.Model, "gpt-4o")
|
||||
}
|
||||
|
||||
if len(result.Input) != 1 {
|
||||
t.Errorf("Input length = %d, want 1", len(result.Input))
|
||||
}
|
||||
|
||||
if result.Input[0].Role != "user" {
|
||||
t.Errorf("Input[0].Role = %q, want %q", result.Input[0].Role, "user")
|
||||
}
|
||||
|
||||
if result.Instructions != "You are a helpful assistant." {
|
||||
t.Errorf("Instructions = %q, want %q", result.Instructions, "You are a helpful assistant.")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesToChat(t *testing.T) {
|
||||
resp := &ResponsesResponse{
|
||||
ID: "resp-123",
|
||||
Model: "gpt-4o",
|
||||
Status: "completed",
|
||||
Output: []OutputItem{
|
||||
{
|
||||
Type: "message",
|
||||
Content: []OutputContent{
|
||||
{Type: "output_text", Text: "2+2 equals 4."},
|
||||
},
|
||||
},
|
||||
},
|
||||
Usage: Usage{
|
||||
PromptTokens: 15,
|
||||
CompletionTokens: 10,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := ResponsesToChat(resp)
|
||||
if err != nil {
|
||||
t.Fatalf("ResponsesToChat() error = %v", err)
|
||||
}
|
||||
|
||||
if result.ID != "resp-123" {
|
||||
t.Errorf("ID = %q, want %q", result.ID, "resp-123")
|
||||
}
|
||||
|
||||
if len(result.Choices) != 1 {
|
||||
t.Errorf("Choices length = %d, want 1", len(result.Choices))
|
||||
return
|
||||
}
|
||||
|
||||
if result.Choices[0].Message.Content != "2+2 equals 4." {
|
||||
t.Errorf("Content = %q, want %q", result.Choices[0].Message.Content, "2+2 equals 4.")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMapStopReason(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
expected string
|
||||
}{
|
||||
{"end_turn", "stop"},
|
||||
{"stop_sequence", "stop"},
|
||||
{"tool_use", "tool_calls"},
|
||||
{"max_tokens", "length"},
|
||||
{"unknown", "stop"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.input, func(t *testing.T) {
|
||||
result := mapStopReason(tt.input)
|
||||
if result != tt.expected {
|
||||
t.Errorf("mapStopReason(%q) = %q, want %q", tt.input, result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMessagesToChatToolUse(t *testing.T) {
|
||||
resp := &MessagesResponse{
|
||||
ID: "msg-456",
|
||||
Model: "claude-3-sonnet-20240229",
|
||||
Content: []ContentBlock{
|
||||
{Type: "text", Text: "Let me search for that."},
|
||||
{Type: "tool_use", ID: "toolu-123", Name: "web_search"},
|
||||
},
|
||||
StopReason: "tool_use",
|
||||
Usage: Usage{
|
||||
PromptTokens: 20,
|
||||
CompletionTokens: 30,
|
||||
},
|
||||
}
|
||||
|
||||
result, err := MessagesToChat(resp)
|
||||
if err != nil {
|
||||
t.Fatalf("MessagesToChat() error = %v", err)
|
||||
}
|
||||
|
||||
if len(result.Choices) != 1 {
|
||||
t.Errorf("Choices length = %d, want 1", len(result.Choices))
|
||||
return
|
||||
}
|
||||
|
||||
if result.Choices[0].FinishReason != "tool_calls" {
|
||||
t.Errorf("FinishReason = %q, want %q", result.Choices[0].FinishReason, "tool_calls")
|
||||
}
|
||||
|
||||
if len(result.Choices[0].Message.ToolCalls) != 1 {
|
||||
t.Errorf("ToolCalls length = %d, want 1", len(result.Choices[0].Message.ToolCalls))
|
||||
return
|
||||
}
|
||||
|
||||
if result.Choices[0].Message.ToolCalls[0].ID != "toolu-123" {
|
||||
t.Errorf("ToolCall ID = %q, want %q", result.Choices[0].Message.ToolCalls[0].ID, "toolu-123")
|
||||
}
|
||||
|
||||
if result.Choices[0].Message.ToolCalls[0].Function.Name != "web_search" {
|
||||
t.Errorf("Function.Name = %q, want %q", result.Choices[0].Message.ToolCalls[0].Function.Name, "web_search")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
package convert
|
||||
|
||||
// MessagesRequest represents an Anthropic Messages API request
|
||||
type MessagesRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []Message `json:"messages"`
|
||||
MaxTokens int `json:"max_tokens"`
|
||||
System interface{} `json:"system,omitempty"` // string or []ContentPart
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty"`
|
||||
TopK *int `json:"top_k,omitempty"`
|
||||
StopSequences []string `json:"stop_sequences,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Tools []Tool `json:"tools,omitempty"`
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
Metadata interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// MessagesResponse represents an Anthropic Messages API response
|
||||
type MessagesResponse struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Content []ContentBlock `json:"content"`
|
||||
Model string `json:"model"`
|
||||
StopReason string `json:"stop_reason"`
|
||||
StopSequence string `json:"stop_sequence,omitempty"`
|
||||
Usage Usage `json:"usage"`
|
||||
}
|
||||
|
||||
// AnthropicStreamEvent represents an Anthropic streaming event
|
||||
type AnthropicStreamEvent struct {
|
||||
Type string `json:"type"`
|
||||
Index int `json:"index,omitempty"`
|
||||
Delta *Delta `json:"delta,omitempty"`
|
||||
Usage *Usage `json:"usage,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package convert
|
||||
|
||||
// ResponsesRequest represents an OpenAI Responses API request
|
||||
type ResponsesRequest struct {
|
||||
Model string `json:"model"`
|
||||
Input []InputItem `json:"input"`
|
||||
Instructions string `json:"instructions,omitempty"`
|
||||
MaxOutputTokens *int `json:"max_output_tokens,omitempty"`
|
||||
Tools []Tool `json:"tools,omitempty"`
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty"`
|
||||
Metadata interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// InputItem represents a single input item
|
||||
type InputItem struct {
|
||||
Role string `json:"role"`
|
||||
Content interface{} `json:"content,omitempty"`
|
||||
}
|
||||
|
||||
// ResponsesResponse represents an OpenAI Responses API response
|
||||
type ResponsesResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
CreatedAt int64 `json:"created_at"`
|
||||
Status string `json:"status"`
|
||||
Model string `json:"model"`
|
||||
Output []OutputItem `json:"output"`
|
||||
Usage Usage `json:"usage"`
|
||||
Error interface{} `json:"error,omitempty"`
|
||||
Incomplete *Incomplete `json:"incomplete,omitempty"`
|
||||
}
|
||||
|
||||
type OutputItem struct {
|
||||
Type string `json:"type"`
|
||||
Content []OutputContent `json:"content,omitempty"`
|
||||
Role string `json:"role,omitempty"`
|
||||
}
|
||||
|
||||
type OutputContent struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Input interface{} `json:"input,omitempty"`
|
||||
}
|
||||
|
||||
type Incomplete struct {
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
// ResponsesStreamEvent represents a Responses API streaming event
|
||||
type ResponsesStreamEvent struct {
|
||||
Type string `json:"type"`
|
||||
Item *OutputItem `json:"item,omitempty"`
|
||||
Delta string `json:"delta,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
package convert
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// SSEWriter writes Server-Sent Events
|
||||
type SSEWriter struct {
|
||||
writer io.Writer
|
||||
flusher http.Flusher
|
||||
}
|
||||
|
||||
// NewSSEWriter creates a new SSE writer
|
||||
func NewSSEWriter(w http.ResponseWriter) *SSEWriter {
|
||||
flusher, _ := w.(http.Flusher)
|
||||
return &SSEWriter{
|
||||
writer: w,
|
||||
flusher: flusher,
|
||||
}
|
||||
}
|
||||
|
||||
// WriteEvent writes a single SSE event
|
||||
func (w *SSEWriter) WriteEvent(event string, data interface{}) error {
|
||||
var dataStr string
|
||||
switch v := data.(type) {
|
||||
case string:
|
||||
dataStr = v
|
||||
default:
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
dataStr = string(b)
|
||||
}
|
||||
|
||||
_, err := fmt.Fprintf(w.writer, "event: %s\ndata: %s\n\n", event, dataStr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if w.flusher != nil {
|
||||
w.flusher.Flush()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// WriteChunk writes a streaming chunk in SSE format
|
||||
func (w *SSEWriter) WriteChunk(chunk interface{}) error {
|
||||
b, err := json.Marshal(chunk)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = fmt.Fprintf(w.writer, "data: %s\n\n", string(b))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if w.flusher != nil {
|
||||
w.flusher.Flush()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// WriteDone writes the [DONE] marker
|
||||
func (w *SSEWriter) WriteDone() error {
|
||||
_, err := fmt.Fprintf(w.writer, "data: [DONE]\n\n")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if w.flusher != nil {
|
||||
w.flusher.Flush()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SSEParser parses Server-Sent Events from a reader
|
||||
type SSEParser struct {
|
||||
reader *bufio.Reader
|
||||
}
|
||||
|
||||
// NewSSEParser creates a new SSE parser
|
||||
func NewSSEParser(r io.Reader) *SSEParser {
|
||||
return &SSEParser{
|
||||
reader: bufio.NewReader(r),
|
||||
}
|
||||
}
|
||||
|
||||
// SSEEvent represents a parsed SSE event
|
||||
type SSEEvent struct {
|
||||
Event string
|
||||
Data string
|
||||
}
|
||||
|
||||
// ReadEvent reads the next SSE event
|
||||
func (p *SSEParser) ReadEvent() (*SSEEvent, error) {
|
||||
event := &SSEEvent{}
|
||||
|
||||
for {
|
||||
line, err := p.reader.ReadString('\n')
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
|
||||
if line == "" {
|
||||
// Empty line means end of event
|
||||
if event.Data != "" || event.Event != "" {
|
||||
return event, nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(line, "event:") {
|
||||
event.Event = strings.TrimSpace(line[6:])
|
||||
} else if strings.HasPrefix(line, "data:") {
|
||||
data := strings.TrimSpace(line[5:])
|
||||
if event.Data != "" {
|
||||
event.Data += "\n" + data
|
||||
} else {
|
||||
event.Data = data
|
||||
}
|
||||
}
|
||||
// Ignore comments (lines starting with :) and unknown fields
|
||||
}
|
||||
}
|
||||
|
||||
// ParseChatStreamChunk parses an OpenAI Chat Completions streaming chunk
|
||||
func ParseChatStreamChunk(data string) (*ChatCompletionStreamChunk, error) {
|
||||
if data == "[DONE]" {
|
||||
return nil, io.EOF
|
||||
}
|
||||
|
||||
var chunk ChatCompletionStreamChunk
|
||||
err := json.Unmarshal([]byte(data), &chunk)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &chunk, nil
|
||||
}
|
||||
|
||||
// ParseMessagesStreamEvent parses an Anthropic Messages streaming event
|
||||
func ParseMessagesStreamEvent(data string) (*AnthropicStreamEvent, error) {
|
||||
var event AnthropicStreamEvent
|
||||
err := json.Unmarshal([]byte(data), &event)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &event, nil
|
||||
}
|
||||
|
||||
// ParseResponsesStreamChunk parses an OpenAI Responses API streaming chunk
|
||||
func ParseResponsesStreamChunk(data string) (*ResponsesStreamEvent, error) {
|
||||
if data == "[DONE]" {
|
||||
return nil, io.EOF
|
||||
}
|
||||
|
||||
var event ResponsesStreamEvent
|
||||
err := json.Unmarshal([]byte(data), &event)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &event, nil
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package convert
|
||||
|
||||
// Common types shared across all protocols
|
||||
|
||||
// Message represents a unified message format
|
||||
type Message struct {
|
||||
Role string `json:"role"`
|
||||
Content interface{} `json:"content,omitempty"` // string or []ContentPart
|
||||
Name string `json:"name,omitempty"`
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
}
|
||||
|
||||
// ContentPart represents a part of a multi-part message content
|
||||
type ContentPart struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
ImageURL *ImageURL `json:"image_url,omitempty"`
|
||||
Source *ImageSource `json:"source,omitempty"`
|
||||
ToolUse *ToolUse `json:"tool_use,omitempty"`
|
||||
ToolResult *ToolResult `json:"tool_result,omitempty"`
|
||||
}
|
||||
|
||||
type ImageURL struct {
|
||||
URL string `json:"url"`
|
||||
Detail string `json:"detail,omitempty"`
|
||||
}
|
||||
|
||||
type ImageSource struct {
|
||||
Type string `json:"type"`
|
||||
MediaType string `json:"media_type"`
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
type ToolCall struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function FunctionCall `json:"function"`
|
||||
}
|
||||
|
||||
type FunctionCall struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
}
|
||||
|
||||
type ToolUse struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Input interface{} `json:"input"`
|
||||
}
|
||||
|
||||
type ToolResult struct {
|
||||
ToolUseID string `json:"tool_use_id"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// Tool definition
|
||||
type Tool struct {
|
||||
Type string `json:"type"`
|
||||
Function ToolDefinition `json:"function,omitempty"`
|
||||
Name string `json:"name,omitempty"` // Anthropic style
|
||||
Input interface{} `json:"input_schema,omitempty"` // Anthropic style
|
||||
}
|
||||
|
||||
type ToolDefinition struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Parameters interface{} `json:"parameters,omitempty"`
|
||||
}
|
||||
|
||||
// StreamEvent represents a unified streaming event
|
||||
type StreamEvent struct {
|
||||
Type string `json:"type"` // "message_start", "content_block_start", "content_block_delta", "message_delta", "message_stop"
|
||||
Delta *Delta `json:"delta,omitempty"`
|
||||
Usage *Usage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
type Delta struct {
|
||||
Type string `json:"type,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
StopReason string `json:"stop_reason,omitempty"`
|
||||
ContentBlock *ContentBlock `json:"content_block,omitempty"`
|
||||
}
|
||||
|
||||
type ContentBlock struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Input interface{} `json:"input,omitempty"`
|
||||
}
|
||||
|
||||
// Usage represents token usage
|
||||
type Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens,omitempty"`
|
||||
CacheReadTokens int `json:"cache_read_input_tokens,omitempty"`
|
||||
}
|
||||
@@ -0,0 +1,380 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"opencatd-open/internal/channel"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/proxy/convert"
|
||||
"opencatd-open/internal/store"
|
||||
"opencatd-open/pkg/config"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type Gateway struct {
|
||||
ctx context.Context
|
||||
cfg *config.Config
|
||||
db *gorm.DB
|
||||
wg *sync.WaitGroup
|
||||
httpClient *http.Client
|
||||
|
||||
userDAO *dao.UserDAO
|
||||
apiKeyDAO *dao.ApiKeyDAO
|
||||
usageDAO *dao.UsageDAO
|
||||
dailyDAO *dao.DailyUsageDAO
|
||||
channelSvc *channel.Service
|
||||
}
|
||||
|
||||
func NewGateway(ctx context.Context, cfg *config.Config, db *gorm.DB, wg *sync.WaitGroup, userDAO *dao.UserDAO, apiKeyDAO *dao.ApiKeyDAO, usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Gateway {
|
||||
client := &http.Client{Timeout: 120 * time.Second}
|
||||
if os.Getenv("LOCAL_PROXY") != "" {
|
||||
proxyUrl, err := url.Parse(os.Getenv("LOCAL_PROXY"))
|
||||
if err == nil {
|
||||
tr := &http.Transport{
|
||||
Proxy: http.ProxyURL(proxyUrl),
|
||||
}
|
||||
client.Transport = tr
|
||||
}
|
||||
}
|
||||
|
||||
return &Gateway{
|
||||
ctx: ctx,
|
||||
cfg: cfg,
|
||||
db: db,
|
||||
wg: wg,
|
||||
httpClient: client,
|
||||
userDAO: userDAO,
|
||||
apiKeyDAO: apiKeyDAO,
|
||||
usageDAO: usageDAO,
|
||||
dailyDAO: dailyDAO,
|
||||
channelSvc: nil,
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) SetChannelService(svc *channel.Service) {
|
||||
g.channelSvc = svc
|
||||
}
|
||||
|
||||
// Request represents a parsed incoming request
|
||||
type Request struct {
|
||||
Model string
|
||||
Stream bool
|
||||
Protocol string // "chat", "messages", "responses"
|
||||
Body []byte
|
||||
APIKey *store.APIKey
|
||||
UserID uint64
|
||||
}
|
||||
|
||||
// ParseRequest parses the incoming request and extracts key fields
|
||||
func (g *Gateway) ParseRequest(c *gin.Context, protocol string) (*Request, error) {
|
||||
body, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read body: %w", err)
|
||||
}
|
||||
|
||||
apiKey, _ := c.Get("api_key")
|
||||
userID, _ := c.Get("user_id")
|
||||
|
||||
req := &Request{
|
||||
Protocol: protocol,
|
||||
Body: body,
|
||||
UserID: userID.(uint64),
|
||||
}
|
||||
|
||||
if ak, ok := apiKey.(*store.APIKey); ok {
|
||||
req.APIKey = ak
|
||||
}
|
||||
|
||||
// Parse model and stream based on protocol
|
||||
switch protocol {
|
||||
case "chat":
|
||||
var parsed convert.ChatCompletionRequest
|
||||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("invalid chat request: %w", err)
|
||||
}
|
||||
req.Model = parsed.Model
|
||||
req.Stream = parsed.Stream
|
||||
case "messages":
|
||||
var parsed convert.MessagesRequest
|
||||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("invalid messages request: %w", err)
|
||||
}
|
||||
req.Model = parsed.Model
|
||||
req.Stream = parsed.Stream
|
||||
case "responses":
|
||||
var parsed convert.ResponsesRequest
|
||||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("invalid responses request: %w", err)
|
||||
}
|
||||
req.Model = parsed.Model
|
||||
req.Stream = parsed.Stream
|
||||
}
|
||||
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// Dispatch routes the request to the appropriate upstream
|
||||
func (g *Gateway) Dispatch(c *gin.Context, req *Request) {
|
||||
if g.channelSvc == nil {
|
||||
g.writeError(c, http.StatusBadGateway, "channel service not available")
|
||||
return
|
||||
}
|
||||
|
||||
ch, err := g.channelSvc.SelectChannel(g.ctx, req.Model)
|
||||
if err != nil {
|
||||
g.writeError(c, http.StatusBadGateway, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
apiKey, err := g.channelSvc.GetAPIKey(ch)
|
||||
if err != nil {
|
||||
g.writeError(c, http.StatusBadGateway, "failed to decrypt API key")
|
||||
return
|
||||
}
|
||||
|
||||
// Determine target format and convert if needed
|
||||
targetFormat := req.Protocol
|
||||
if len(ch.FormatsEffective()) > 0 {
|
||||
// Prefer the channel's native format
|
||||
for _, f := range ch.FormatsEffective() {
|
||||
if f == req.Protocol {
|
||||
targetFormat = f
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build upstream URL
|
||||
upstreamPath := g.getUpstreamPath(req.Protocol)
|
||||
upstreamURL := ch.UpstreamURL(req.Protocol, upstreamPath)
|
||||
|
||||
// Convert request if needed
|
||||
var requestBody []byte
|
||||
if targetFormat != req.Protocol {
|
||||
requestBody, err = g.convertRequest(req.Body, req.Protocol, targetFormat)
|
||||
if err != nil {
|
||||
g.writeError(c, http.StatusBadRequest, "conversion failed: "+err.Error())
|
||||
return
|
||||
}
|
||||
} else {
|
||||
requestBody = req.Body
|
||||
}
|
||||
|
||||
// Create upstream request
|
||||
httpReq, err := http.NewRequestWithContext(g.ctx, "POST", upstreamURL, bytes.NewReader(requestBody))
|
||||
if err != nil {
|
||||
g.writeError(c, http.StatusBadGateway, "failed to create request")
|
||||
return
|
||||
}
|
||||
|
||||
// Set headers
|
||||
g.setHeaders(httpReq, ch, apiKey, targetFormat)
|
||||
|
||||
// Execute request
|
||||
start := time.Now()
|
||||
resp, err := g.httpClient.Do(httpReq)
|
||||
latency := time.Since(start)
|
||||
if err != nil {
|
||||
g.channelSvc.RecordFailure(ch.ID)
|
||||
g.writeError(c, http.StatusBadGateway, fmt.Sprintf("upstream error: %v (latency: %v)", err, latency))
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Record success
|
||||
g.channelSvc.RecordSuccess(ch.ID)
|
||||
|
||||
// Handle response
|
||||
if resp.StatusCode >= 400 {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
log.Printf("Upstream error: status=%d body=%s", resp.StatusCode, string(body))
|
||||
c.Data(resp.StatusCode, "application/json", body)
|
||||
return
|
||||
}
|
||||
|
||||
// Stream or buffer response
|
||||
if req.Stream {
|
||||
g.streamResponse(c, resp, req.Protocol, ch)
|
||||
} else {
|
||||
g.bufferResponse(c, resp, req.Protocol, ch)
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) getUpstreamPath(protocol string) string {
|
||||
switch protocol {
|
||||
case "chat":
|
||||
return "/chat/completions"
|
||||
case "messages":
|
||||
return "/messages"
|
||||
case "responses":
|
||||
return "/responses"
|
||||
default:
|
||||
return "/chat/completions"
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) setHeaders(req *http.Request, ch *store.Channel, apiKey string, format string) {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
switch ch.Provider {
|
||||
case store.ChannelProviderOpenAI, store.ChannelProviderCompatible:
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
case store.ChannelProviderAnthropic:
|
||||
req.Header.Set("x-api-key", apiKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) convertRequest(body []byte, from, to string) ([]byte, error) {
|
||||
switch {
|
||||
case from == "chat" && to == "messages":
|
||||
var req convert.ChatCompletionRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
msgReq, err := convert.ChatToMessages(&req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return json.Marshal(msgReq)
|
||||
|
||||
case from == "chat" && to == "responses":
|
||||
var req convert.ChatCompletionRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
respReq, err := convert.ChatToResponses(&req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return json.Marshal(respReq)
|
||||
|
||||
case from == "messages" && to == "chat":
|
||||
var req convert.MessagesRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Messages -> Chat: we need to construct a ChatCompletionRequest
|
||||
chatReq := &convert.ChatCompletionRequest{
|
||||
Model: req.Model,
|
||||
}
|
||||
for _, m := range req.Messages {
|
||||
chatReq.Messages = append(chatReq.Messages, m)
|
||||
}
|
||||
if req.Temperature != nil {
|
||||
chatReq.Temperature = req.Temperature
|
||||
}
|
||||
if req.TopP != nil {
|
||||
chatReq.TopP = req.TopP
|
||||
}
|
||||
chatReq.Tools = req.Tools
|
||||
chatReq.Stream = req.Stream
|
||||
return json.Marshal(chatReq)
|
||||
|
||||
case from == "responses" && to == "chat":
|
||||
var req convert.ResponsesRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chatReq := &convert.ChatCompletionRequest{
|
||||
Model: req.Model,
|
||||
}
|
||||
for _, item := range req.Input {
|
||||
chatReq.Messages = append(chatReq.Messages, convert.Message{
|
||||
Role: item.Role,
|
||||
Content: item.Content,
|
||||
})
|
||||
}
|
||||
chatReq.Tools = req.Tools
|
||||
chatReq.Stream = req.Stream
|
||||
return json.Marshal(chatReq)
|
||||
|
||||
default:
|
||||
return body, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) streamResponse(c *gin.Context, resp *http.Response, protocol string, ch *store.Channel) {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Status(http.StatusOK)
|
||||
|
||||
writer := convert.NewSSEWriter(c.Writer)
|
||||
parser := convert.NewSSEParser(resp.Body)
|
||||
|
||||
for {
|
||||
event, err := parser.ReadEvent()
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
log.Printf("Stream parse error: %v", err)
|
||||
break
|
||||
}
|
||||
|
||||
if event.Event == "error" {
|
||||
log.Printf("Upstream stream error: %s", event.Data)
|
||||
break
|
||||
}
|
||||
|
||||
// Write raw SSE event based on protocol
|
||||
if err := writer.WriteEvent("chat CompletionChunk", event.Data); err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
writer.WriteDone()
|
||||
}
|
||||
|
||||
func (g *Gateway) bufferResponse(c *gin.Context, resp *http.Response, protocol string, ch *store.Channel) {
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
g.writeError(c, http.StatusBadGateway, "failed to read response")
|
||||
return
|
||||
}
|
||||
|
||||
c.Data(resp.StatusCode, "application/json", body)
|
||||
}
|
||||
|
||||
func (g *Gateway) writeError(c *gin.Context, status int, message string) {
|
||||
protocol := c.GetHeader("X-Protocol")
|
||||
if protocol == "" {
|
||||
protocol = "chat"
|
||||
}
|
||||
|
||||
switch {
|
||||
case strings.Contains(c.GetHeader("Accept"), "text/event-stream"):
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Status(status)
|
||||
fmt.Fprintf(c.Writer, "data: {\"error\":{\"message\":\"%s\"}}\n\n", message)
|
||||
fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
|
||||
case protocol == "messages":
|
||||
c.JSON(status, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{
|
||||
"type": "api_error",
|
||||
"message": message,
|
||||
},
|
||||
})
|
||||
default:
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"message": message,
|
||||
"type": "invalid_request_error",
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// HandleChat handles POST /v1/chat/completions
|
||||
func (g *Gateway) HandleChat(c *gin.Context) {
|
||||
req, err := g.ParseRequest(c, "chat")
|
||||
if err != nil {
|
||||
g.writeError(c, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
g.Dispatch(c, req)
|
||||
}
|
||||
|
||||
// HandleMessages handles POST /v1/messages
|
||||
func (g *Gateway) HandleMessages(c *gin.Context) {
|
||||
req, err := g.ParseRequest(c, "messages")
|
||||
if err != nil {
|
||||
g.writeError(c, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
g.Dispatch(c, req)
|
||||
}
|
||||
|
||||
// HandleResponses handles POST /v1/responses
|
||||
func (g *Gateway) HandleResponses(c *gin.Context) {
|
||||
req, err := g.ParseRequest(c, "responses")
|
||||
if err != nil {
|
||||
g.writeError(c, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
g.Dispatch(c, req)
|
||||
}
|
||||
|
||||
// HandleModels handles GET /v1/models
|
||||
func (g *Gateway) HandleModels(c *gin.Context) {
|
||||
// TODO: Return list of available models based on enabled channels
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": []interface{}{},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type ApiKeyServiceImpl struct {
|
||||
db *gorm.DB
|
||||
apiKeyRepo *dao.ApiKeyDAO
|
||||
}
|
||||
|
||||
func NewApiKeyService(db *gorm.DB, apiKeyDao *dao.ApiKeyDAO) *ApiKeyServiceImpl {
|
||||
return &ApiKeyServiceImpl{db: db, apiKeyRepo: apiKeyDao}
|
||||
}
|
||||
|
||||
func (s *ApiKeyServiceImpl) CreateApiKey(ctx context.Context, apikey *store.APIKey) error {
|
||||
return s.apiKeyRepo.Create(apikey)
|
||||
}
|
||||
|
||||
func (s *ApiKeyServiceImpl) GetApiKey(ctx context.Context, id uint64) (*store.APIKey, error) {
|
||||
return s.apiKeyRepo.GetByID(id)
|
||||
}
|
||||
|
||||
func (s *ApiKeyServiceImpl) ListApiKey(ctx context.Context, userID uint64, limit, offset int) ([]*store.APIKey, int64, error) {
|
||||
return s.apiKeyRepo.ListByUserID(userID, limit, offset)
|
||||
}
|
||||
|
||||
func (s *ApiKeyServiceImpl) UpdateApiKey(ctx context.Context, apikey *store.APIKey) error {
|
||||
return s.apiKeyRepo.Update(apikey)
|
||||
}
|
||||
|
||||
func (s *ApiKeyServiceImpl) DeleteApiKey(ctx context.Context, id uint64) error {
|
||||
return s.apiKeyRepo.Delete(id)
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"opencatd-open/internal/channel"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
"opencatd-open/internal/pkg/crypto"
|
||||
)
|
||||
|
||||
type ChannelServiceImpl struct {
|
||||
channelDAO *dao.ChannelDAO
|
||||
channelSvc *channel.Service
|
||||
}
|
||||
|
||||
func NewChannelService(channelDAO *dao.ChannelDAO, channelSvc *channel.Service) *ChannelServiceImpl {
|
||||
return &ChannelServiceImpl{
|
||||
channelDAO: channelDAO,
|
||||
channelSvc: channelSvc,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ChannelServiceImpl) Create(ctx context.Context, ch *store.Channel) error {
|
||||
return s.channelDAO.Create(ch)
|
||||
}
|
||||
|
||||
func (s *ChannelServiceImpl) GetByID(ctx context.Context, id uint64) (*store.Channel, error) {
|
||||
return s.channelDAO.GetByID(id)
|
||||
}
|
||||
|
||||
func (s *ChannelServiceImpl) List(ctx context.Context, limit, offset int) ([]*store.Channel, int64, error) {
|
||||
return s.channelDAO.List(limit, offset)
|
||||
}
|
||||
|
||||
func (s *ChannelServiceImpl) ListEnabled(ctx context.Context) ([]*store.Channel, error) {
|
||||
return s.channelDAO.ListEnabled()
|
||||
}
|
||||
|
||||
func (s *ChannelServiceImpl) Update(ctx context.Context, ch *store.Channel) error {
|
||||
return s.channelDAO.Update(ch)
|
||||
}
|
||||
|
||||
func (s *ChannelServiceImpl) Delete(ctx context.Context, id uint64) error {
|
||||
return s.channelDAO.Delete(id)
|
||||
}
|
||||
|
||||
// GetAPIKey decrypts the channel's API key
|
||||
func (s *ChannelServiceImpl) GetAPIKey(ctx context.Context, channelID uint64) (string, error) {
|
||||
ch, err := s.channelDAO.GetByID(channelID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return crypto.Decrypt(ch.APIKeyEnc)
|
||||
}
|
||||
|
||||
// SelectForModel selects the best channel for a model
|
||||
func (s *ChannelServiceImpl) SelectForModel(ctx context.Context, modelName string) (*store.Channel, error) {
|
||||
return s.channelSvc.SelectChannel(ctx, modelName)
|
||||
}
|
||||
|
||||
// BindModels binds models to a channel
|
||||
func (s *ChannelServiceImpl) BindModels(ctx context.Context, channelID uint64, bindings []store.ChannelModelBinding) error {
|
||||
return s.channelDAO.BindModels(channelID, bindings)
|
||||
}
|
||||
|
||||
// GetChannelModels returns models bound to a channel
|
||||
func (s *ChannelServiceImpl) GetChannelModels(ctx context.Context, channelID uint64) ([]store.ChannelModelBinding, error) {
|
||||
return s.channelDAO.GetChannelModels(channelID)
|
||||
}
|
||||
|
||||
// GetModelChannels returns channels for a model
|
||||
func (s *ChannelServiceImpl) GetModelChannels(ctx context.Context, modelName string) ([]*store.Channel, error) {
|
||||
return s.channelDAO.GetEnabledChannelsByModel(modelName)
|
||||
}
|
||||
|
||||
// RecordSuccess records a successful request
|
||||
func (s *ChannelServiceImpl) RecordSuccess(channelID uint64) {
|
||||
s.channelSvc.RecordSuccess(channelID)
|
||||
}
|
||||
|
||||
// RecordFailure records a failed request
|
||||
func (s *ChannelServiceImpl) RecordFailure(channelID uint64) {
|
||||
s.channelSvc.RecordFailure(channelID)
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
)
|
||||
|
||||
type ModelServiceImpl struct {
|
||||
modelDAO *dao.ModelDAO
|
||||
channelDAO *dao.ChannelDAO
|
||||
}
|
||||
|
||||
func NewModelService(modelDAO *dao.ModelDAO, channelDAO *dao.ChannelDAO) *ModelServiceImpl {
|
||||
return &ModelServiceImpl{
|
||||
modelDAO: modelDAO,
|
||||
channelDAO: channelDAO,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ModelServiceImpl) Create(ctx context.Context, model *store.Model) error {
|
||||
return s.modelDAO.Create(model)
|
||||
}
|
||||
|
||||
func (s *ModelServiceImpl) GetByID(ctx context.Context, id uint64) (*store.Model, error) {
|
||||
return s.modelDAO.GetByID(id)
|
||||
}
|
||||
|
||||
func (s *ModelServiceImpl) GetByName(ctx context.Context, name string) (*store.Model, error) {
|
||||
return s.modelDAO.GetByName(name)
|
||||
}
|
||||
|
||||
func (s *ModelServiceImpl) List(ctx context.Context, limit, offset int) ([]*store.Model, int64, error) {
|
||||
return s.modelDAO.List(limit, offset)
|
||||
}
|
||||
|
||||
func (s *ModelServiceImpl) ListEnabled(ctx context.Context) ([]*store.Model, error) {
|
||||
return s.modelDAO.ListEnabled()
|
||||
}
|
||||
|
||||
func (s *ModelServiceImpl) Update(ctx context.Context, model *store.Model) error {
|
||||
return s.modelDAO.Update(model)
|
||||
}
|
||||
|
||||
func (s *ModelServiceImpl) Delete(ctx context.Context, id uint64) error {
|
||||
return s.modelDAO.Delete(id)
|
||||
}
|
||||
|
||||
func (s *ModelServiceImpl) Upsert(ctx context.Context, model *store.Model) error {
|
||||
return s.modelDAO.Upsert(model)
|
||||
}
|
||||
|
||||
// BindChannel binds a model to a channel
|
||||
func (s *ModelServiceImpl) BindChannel(ctx context.Context, modelID, channelID uint64, upstreamModel string, weight int) error {
|
||||
binding := store.ChannelModelBinding{
|
||||
ModelID: modelID,
|
||||
ChannelID: channelID,
|
||||
UpstreamModel: upstreamModel,
|
||||
Weight: weight,
|
||||
}
|
||||
return s.channelDAO.BindModels(channelID, []store.ChannelModelBinding{binding})
|
||||
}
|
||||
|
||||
// ListChannelModels lists all models bound to a channel
|
||||
func (s *ModelServiceImpl) ListChannelModels(ctx context.Context, channelID uint64) ([]store.ChannelModelBinding, error) {
|
||||
return s.channelDAO.GetChannelModels(channelID)
|
||||
}
|
||||
|
||||
// ListModelChannels lists all channels for a model
|
||||
func (s *ModelServiceImpl) ListModelChannels(ctx context.Context, modelName string) ([]*store.Channel, error) {
|
||||
return s.channelDAO.GetEnabledChannelsByModel(modelName)
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type TokenServiceImpl struct {
|
||||
db *gorm.DB
|
||||
tokenRepo *dao.TokenDAO
|
||||
}
|
||||
|
||||
func NewTokenService(db *gorm.DB, tokenRepo *dao.TokenDAO) *TokenServiceImpl {
|
||||
return &TokenServiceImpl{
|
||||
db: db,
|
||||
tokenRepo: tokenRepo,
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TokenServiceImpl) GetByKey(ctx context.Context, key string) (*store.User, error) {
|
||||
return t.tokenRepo.GetByKey(key)
|
||||
}
|
||||
|
||||
func (t *TokenServiceImpl) GetByID(ctx context.Context, id uint64) (*store.User, error) {
|
||||
return t.tokenRepo.GetByID(id)
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"opencatd-open/pkg/config"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type UsageService struct {
|
||||
Ctx context.Context
|
||||
Cfg *config.Config
|
||||
DB *gorm.DB
|
||||
}
|
||||
|
||||
func NewUsageService(ctx context.Context, cfg *config.Config, db *gorm.DB) *UsageService {
|
||||
return &UsageService{
|
||||
Ctx: ctx,
|
||||
Cfg: cfg,
|
||||
DB: db,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
"opencatd-open/pkg/config"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type UserServiceImpl struct {
|
||||
cfg *config.Config
|
||||
db *gorm.DB
|
||||
userRepo *dao.UserDAO
|
||||
}
|
||||
|
||||
func NewUserService(cfg *config.Config, db *gorm.DB, userRepo *dao.UserDAO) *UserServiceImpl {
|
||||
return &UserServiceImpl{
|
||||
cfg: cfg,
|
||||
db: db,
|
||||
userRepo: userRepo,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *UserServiceImpl) GetByID(ctx context.Context, id uint64) (*store.User, error) {
|
||||
return s.userRepo.GetByID(id)
|
||||
}
|
||||
|
||||
func (s *UserServiceImpl) GetByUsername(ctx context.Context, username string) (*store.User, error) {
|
||||
return s.userRepo.GetByUsername(username)
|
||||
}
|
||||
|
||||
func (s *UserServiceImpl) List(ctx context.Context, limit, offset int) ([]*store.User, int64, error) {
|
||||
return s.userRepo.List(limit, offset)
|
||||
}
|
||||
|
||||
func (s *UserServiceImpl) Create(ctx context.Context, user *store.User) error {
|
||||
return s.userRepo.Create(user)
|
||||
}
|
||||
|
||||
func (s *UserServiceImpl) Update(ctx context.Context, user *store.User) error {
|
||||
return s.userRepo.Update(user)
|
||||
}
|
||||
|
||||
func (s *UserServiceImpl) Delete(ctx context.Context, id uint64) error {
|
||||
return s.userRepo.Delete(id)
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"opencatd-open/internal/store"
|
||||
"opencatd-open/pkg/config"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/go-webauthn/webauthn/protocol"
|
||||
"github.com/go-webauthn/webauthn/webauthn"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type WebAuthnUser struct {
|
||||
User *store.User
|
||||
Credentials []webauthn.Credential
|
||||
}
|
||||
|
||||
func (u *WebAuthnUser) WebAuthnID() []byte {
|
||||
return []byte(strconv.FormatUint(u.User.ID, 10))
|
||||
}
|
||||
|
||||
func (u *WebAuthnUser) WebAuthnName() string {
|
||||
return u.User.Username
|
||||
}
|
||||
|
||||
func (u *WebAuthnUser) WebAuthnDisplayName() string {
|
||||
return u.User.Username
|
||||
}
|
||||
|
||||
func (u *WebAuthnUser) WebAuthnCredentials() []webauthn.Credential {
|
||||
return u.Credentials
|
||||
}
|
||||
|
||||
func (u *WebAuthnUser) WebAuthnCredentialDescriptors() (descriptors []protocol.CredentialDescriptor) {
|
||||
credentials := u.WebAuthnCredentials()
|
||||
descriptors = make([]protocol.CredentialDescriptor, len(credentials))
|
||||
for i, credential := range credentials {
|
||||
descriptors[i] = credential.Descriptor()
|
||||
}
|
||||
return descriptors
|
||||
}
|
||||
|
||||
type WebAuthnService struct {
|
||||
cfg *config.Config
|
||||
DB *gorm.DB
|
||||
WebAuthn *webauthn.WebAuthn
|
||||
}
|
||||
|
||||
func NewWebAuthnService(cfg *config.Config, db *gorm.DB) (*WebAuthnService, error) {
|
||||
wconfig := &webauthn.Config{
|
||||
RPDisplayName: cfg.AppName,
|
||||
RPID: cfg.RPID,
|
||||
RPOrigins: cfg.RPOrigins,
|
||||
AuthenticatorSelection: protocol.AuthenticatorSelection{
|
||||
RequireResidentKey: protocol.ResidentKeyRequired(),
|
||||
ResidentKey: protocol.ResidentKeyRequirementRequired,
|
||||
UserVerification: protocol.VerificationPreferred,
|
||||
},
|
||||
}
|
||||
|
||||
wa, err := webauthn.New(wconfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &WebAuthnService{
|
||||
cfg: cfg,
|
||||
DB: db,
|
||||
WebAuthn: wa,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *WebAuthnService) GetUserWithCredentials(userID uint64) (*WebAuthnUser, error) {
|
||||
var user store.User
|
||||
if err := s.DB.First(&user, userID).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var passkeys []store.Passkey
|
||||
if err := s.DB.Where("user_id = ?", userID).Find(&passkeys).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
credentials := make([]webauthn.Credential, len(passkeys))
|
||||
for i, pk := range passkeys {
|
||||
credentialIDBytes, err := base64.StdEncoding.DecodeString(pk.CredentialID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to decode CredentialID: %w", err)
|
||||
}
|
||||
publicKeyBytes, err := base64.StdEncoding.DecodeString(pk.PublicKey)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to decode PublicKey: %w", err)
|
||||
}
|
||||
aaguidBytes, err := base64.StdEncoding.DecodeString(pk.AAGUID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to decode AAGUID: %w", err)
|
||||
}
|
||||
|
||||
var transport []protocol.AuthenticatorTransport
|
||||
if pk.Transport != "" {
|
||||
transport = []protocol.AuthenticatorTransport{protocol.AuthenticatorTransport(pk.Transport)}
|
||||
}
|
||||
|
||||
credentials[i] = webauthn.Credential{
|
||||
ID: credentialIDBytes,
|
||||
PublicKey: publicKeyBytes,
|
||||
AttestationType: pk.AttestationType,
|
||||
Transport: transport,
|
||||
Flags: webauthn.CredentialFlags{
|
||||
UserPresent: true,
|
||||
UserVerified: true,
|
||||
BackupEligible: pk.BackupEligible,
|
||||
BackupState: pk.BackupState,
|
||||
},
|
||||
Authenticator: webauthn.Authenticator{
|
||||
AAGUID: aaguidBytes,
|
||||
SignCount: uint32(pk.SignCount),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
return &WebAuthnUser{
|
||||
User: &user,
|
||||
Credentials: credentials,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *WebAuthnService) BeginRegistration(userID uint64) (*protocol.CredentialCreation, error) {
|
||||
user, err := s.GetUserWithCredentials(userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
options, _, err := s.WebAuthn.BeginRegistration(user)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return options, nil
|
||||
}
|
||||
|
||||
func (s *WebAuthnService) FinishRegistration(userID uint64, response *http.Request, deviceName string) (*store.Passkey, error) {
|
||||
user, err := s.GetUserWithCredentials(userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
credential, err := s.WebAuthn.FinishRegistration(user, webauthn.SessionData{}, response)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var transport string
|
||||
if len(credential.Transport) > 0 {
|
||||
transport = string(credential.Transport[0])
|
||||
}
|
||||
|
||||
passkey := &store.Passkey{
|
||||
UserID: userID,
|
||||
CredentialID: base64.StdEncoding.EncodeToString(credential.ID),
|
||||
PublicKey: base64.StdEncoding.EncodeToString(credential.PublicKey),
|
||||
AttestationType: string(credential.AttestationType),
|
||||
AAGUID: base64.StdEncoding.EncodeToString(credential.Authenticator.AAGUID),
|
||||
SignCount: uint64(credential.Authenticator.SignCount),
|
||||
Name: deviceName,
|
||||
DeviceType: strings.TrimSpace(fmt.Sprintf("%s", deviceName)),
|
||||
LastUsedAt: time.Now().Unix(),
|
||||
BackupEligible: credential.Flags.BackupEligible,
|
||||
BackupState: credential.Flags.BackupState,
|
||||
Transport: transport,
|
||||
}
|
||||
|
||||
if err := s.DB.Create(passkey).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return passkey, nil
|
||||
}
|
||||
|
||||
func (s *WebAuthnService) BeginLogin() (*protocol.CredentialAssertion, error) {
|
||||
options, _, err := s.WebAuthn.BeginDiscoverableLogin()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return options, nil
|
||||
}
|
||||
|
||||
func (s *WebAuthnService) ListPasskeys(userID uint64) ([]store.Passkey, error) {
|
||||
var passkeys []store.Passkey
|
||||
if err := s.DB.Where("user_id = ?", userID).Find(&passkeys).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return passkeys, nil
|
||||
}
|
||||
|
||||
func (s *WebAuthnService) DeletePasskey(userID uint64, passkeyID uint64) error {
|
||||
return s.DB.Where("id = ? AND user_id = ?", passkeyID, userID).Delete(&store.Passkey{}).Error
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"opencatd-open/pkg/config"
|
||||
|
||||
_ "github.com/lib/pq"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var DB *gorm.DB
|
||||
|
||||
func InitDB(cfg *config.Config) (*gorm.DB, error) {
|
||||
var dialector gorm.Dialector
|
||||
|
||||
switch cfg.DB_Type {
|
||||
case "sqlite":
|
||||
dialector = sqliteDialector(cfg.DSN)
|
||||
case "postgres":
|
||||
dialector = postgresDialector(cfg.DSN)
|
||||
case "mysql":
|
||||
dialector = mysqlDialector(cfg.DSN)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported database type: %s", cfg.DB_Type)
|
||||
}
|
||||
|
||||
db, err := gorm.Open(dialector, &gorm.Config{})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to connect database: %w", err)
|
||||
}
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get underlying *sql.DB: %w", err)
|
||||
}
|
||||
sqlDB.SetMaxOpenConns(cfg.DBMaxOpenConns)
|
||||
sqlDB.SetMaxIdleConns(cfg.DBMaxIdleConns)
|
||||
|
||||
if err := db.AutoMigrate(AllModels()...); err != nil {
|
||||
log.Printf("AutoMigrate warning: %v", err)
|
||||
}
|
||||
|
||||
DB = db
|
||||
return db, nil
|
||||
}
|
||||
|
||||
func sqliteDialector(dsn string) gorm.Dialector {
|
||||
if dsn == "" {
|
||||
dsn = "opencatd.db"
|
||||
}
|
||||
return sqlite.Open(dsn)
|
||||
}
|
||||
|
||||
func postgresDialector(dsn string) gorm.Dialector {
|
||||
if dsn == "" {
|
||||
dsn = "host=localhost user=postgres password=postgres dbname=opencatd port=5432 sslmode=disable"
|
||||
}
|
||||
return postgres.Open(dsn)
|
||||
}
|
||||
|
||||
func mysqlDialector(dsn string) gorm.Dialector {
|
||||
if dsn == "" {
|
||||
dsn = "root:password@tcp(127.0.0.1:3306)/opencatd?charset=utf8mb4&parseTime=True&loc=Local"
|
||||
}
|
||||
return mysql.Open(dsn)
|
||||
}
|
||||
@@ -0,0 +1,233 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 角色 / 状态枚举
|
||||
const (
|
||||
RoleUser = "user"
|
||||
RoleAdmin = "admin"
|
||||
|
||||
UserStatusActive = "active"
|
||||
UserStatusDisabled = "disabled"
|
||||
|
||||
KeyStatusActive = "active"
|
||||
KeyStatusRevoked = "revoked"
|
||||
|
||||
ChannelProviderOpenAI = "openai"
|
||||
ChannelProviderAnthropic = "anthropic"
|
||||
ChannelProviderCompatible = "compatible"
|
||||
ChannelHealthHealthy = "healthy"
|
||||
ChannelHealthDegraded = "degraded"
|
||||
ChannelHealthCooldown = "cooldown"
|
||||
|
||||
FormatChat = "chat"
|
||||
FormatResponses = "responses"
|
||||
FormatMessages = "messages"
|
||||
|
||||
UsageStatusSuccess = "success"
|
||||
UsageStatusError = "error"
|
||||
UsageStatusCanceled = "canceled"
|
||||
)
|
||||
|
||||
// User 用户
|
||||
type User struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
Username string `gorm:"uniqueIndex;size:64;not null" json:"username"`
|
||||
Email string `gorm:"uniqueIndex;size:255;not null" json:"email"`
|
||||
PasswordHash string `gorm:"size:255;not null" json:"-"`
|
||||
Role string `gorm:"size:16;not null;default:user" json:"role"`
|
||||
Balance float64 `gorm:"type:numeric(20,8);not null;default:0" json:"balance"`
|
||||
Status string `gorm:"size:16;not null;default:active" json:"status"`
|
||||
AllowedModels []string `gorm:"type:jsonb;serializer:json" json:"allowed_models,omitempty"`
|
||||
DeniedModels []string `gorm:"type:jsonb;serializer:json" json:"denied_models,omitempty"`
|
||||
InviteCode *string `json:"invite_code,omitempty"`
|
||||
LastLoginAt *time.Time `json:"last_login_at,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// APIKey 密钥(SHA-256 hash 存储)
|
||||
type APIKey struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
UserID uint64 `gorm:"index;not null" json:"user_id"`
|
||||
Name string `gorm:"size:64;not null" json:"name"`
|
||||
KeyHash string `gorm:"uniqueIndex;size:64;not null" json:"-"`
|
||||
KeyPrefix string `gorm:"size:32;not null" json:"key_prefix"`
|
||||
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day,omitempty"`
|
||||
QuotaRequestsPerDay *int `json:"quota_requests_per_day,omitempty"`
|
||||
AllowedModels []string `gorm:"type:jsonb;serializer:json" json:"allowed_models,omitempty"`
|
||||
ExpiresAt *time.Time `json:"expires_at,omitempty"`
|
||||
Status string `gorm:"size:16;not null;default:active" json:"status"`
|
||||
LastUsedAt *time.Time `json:"last_used_at,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Channel 上游渠道
|
||||
type Channel struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
Name string `gorm:"uniqueIndex;size:64;not null" json:"name"`
|
||||
Provider string `gorm:"size:16;not null" json:"provider"`
|
||||
Formats []string `gorm:"type:jsonb;serializer:json" json:"formats,omitempty"`
|
||||
BaseURL string `gorm:"size:255;not null" json:"base_url"`
|
||||
BaseURLs map[string]string `gorm:"type:jsonb;serializer:json" json:"base_urls,omitempty"`
|
||||
APIKeyEnc string `gorm:"size:1024;not null" json:"-"`
|
||||
Weight int `gorm:"not null;default:1" json:"weight"`
|
||||
Priority int `gorm:"not null;default:0" json:"priority"`
|
||||
TimeoutMS int `gorm:"not null;default:120000" json:"timeout_ms"`
|
||||
MaxConcurrency int `gorm:"not null;default:16" json:"max_concurrency"`
|
||||
HealthStatus string `gorm:"size:16;not null;default:healthy" json:"health_status"`
|
||||
Enabled bool `gorm:"not null;default:true" json:"enabled"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// FormatsEffective 返回渠道实际支持的原生协议
|
||||
func (c *Channel) FormatsEffective() []string {
|
||||
if len(c.Formats) > 0 {
|
||||
return c.Formats
|
||||
}
|
||||
switch c.Provider {
|
||||
case ChannelProviderAnthropic:
|
||||
return []string{FormatMessages}
|
||||
case ChannelProviderOpenAI:
|
||||
return []string{FormatChat, FormatResponses}
|
||||
default:
|
||||
return []string{FormatChat}
|
||||
}
|
||||
}
|
||||
|
||||
var versionSegRe = regexp.MustCompile(`/v[0-9]+/?$`)
|
||||
|
||||
// UpstreamURL 按协议选 base_url,拼资源路径
|
||||
func (c *Channel) UpstreamURL(proto, path string) string {
|
||||
base := c.BaseURL
|
||||
if len(c.BaseURLs) > 0 && c.BaseURLs[proto] != "" {
|
||||
base = c.BaseURLs[proto]
|
||||
}
|
||||
base = strings.TrimRight(base, "/")
|
||||
if base == "" {
|
||||
return path
|
||||
}
|
||||
if strings.HasSuffix(base, path) {
|
||||
return base
|
||||
}
|
||||
if versionSegRe.MatchString(base) {
|
||||
return base + path
|
||||
}
|
||||
return base + "/v1" + path
|
||||
}
|
||||
|
||||
// Model 全局模型 + 定价(价格按每百万 token,USD)
|
||||
type Model struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
Name string `gorm:"uniqueIndex;size:128;not null" json:"name"`
|
||||
DisplayName string `gorm:"size:128" json:"display_name"`
|
||||
InputPrice float64 `gorm:"type:numeric(20,8);not null;default:0" json:"input_price"`
|
||||
OutputPrice float64 `gorm:"type:numeric(20,8);not null;default:0" json:"output_price"`
|
||||
CacheReadPrice float64 `gorm:"type:numeric(20,8);not null;default:0" json:"cache_read_price"`
|
||||
Enabled bool `gorm:"not null;default:true" json:"enabled"`
|
||||
Sort int `gorm:"not null;default:0" json:"sort"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// ChannelModelBinding 渠道↔模型绑定(多对多)
|
||||
type ChannelModelBinding struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
ChannelID uint64 `gorm:"index:idx_ch_model,unique;not null" json:"channel_id"`
|
||||
ModelID uint64 `gorm:"index:idx_ch_model,unique;not null" json:"model_id"`
|
||||
UpstreamModel string `gorm:"size:255;not null" json:"upstream_model"`
|
||||
Weight int `gorm:"not null;default:1" json:"weight"`
|
||||
Channel Channel `gorm:"foreignKey:ChannelID" json:"-"`
|
||||
Model Model `gorm:"foreignKey:ModelID" json:"-"`
|
||||
}
|
||||
|
||||
// UsageLog 请求级用量明细
|
||||
type UsageLog struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
RequestID string `gorm:"size:128" json:"request_id"`
|
||||
TraceID string `gorm:"size:64;index" json:"trace_id"`
|
||||
UserID uint64 `gorm:"index:idx_user_created;not null" json:"user_id"`
|
||||
KeyID uint64 `json:"key_id"`
|
||||
ChannelID uint64 `json:"channel_id"`
|
||||
ModelID uint64 `json:"model_id"`
|
||||
ModelName string `gorm:"size:128" json:"model_name"`
|
||||
Protocol string `gorm:"size:32" json:"protocol"`
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
CacheReadTokens int64 `json:"cache_read_tokens"`
|
||||
CacheCreationTokens int64 `json:"cache_creation_tokens"`
|
||||
InputPrice float64 `gorm:"type:numeric(20,8)" json:"input_price"`
|
||||
OutputPrice float64 `gorm:"type:numeric(20,8)" json:"output_price"`
|
||||
CacheReadPrice float64 `gorm:"type:numeric(20,8)" json:"cache_read_price"`
|
||||
Cost float64 `gorm:"type:numeric(20,8)" json:"cost"`
|
||||
LatencyMS int `json:"latency_ms"`
|
||||
Status string `gorm:"size:16;not null" json:"status"`
|
||||
ErrorCode *string `json:"error_code,omitempty"`
|
||||
CreatedAt time.Time `gorm:"index" json:"created_at"`
|
||||
}
|
||||
|
||||
// UsageDaily 日粒度预聚合
|
||||
type UsageDaily struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
UserID uint64 `gorm:"index:idx_daily_user_model,unique" json:"user_id"`
|
||||
ModelID uint64 `gorm:"index:idx_daily_user_model,unique" json:"model_id"`
|
||||
Date string `gorm:"size:10;index:idx_daily_user_model,unique" json:"date"`
|
||||
Requests int64 `json:"requests"`
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
CacheReadTokens int64 `json:"cache_read_tokens"`
|
||||
Cost float64 `gorm:"type:numeric(20,8)" json:"cost"`
|
||||
}
|
||||
|
||||
// Passkey WebAuthn 凭据
|
||||
type Passkey struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
UserID uint64 `gorm:"index;not null" json:"user_id"`
|
||||
Name string `gorm:"size:64" json:"name"`
|
||||
CredentialID string `gorm:"size:255;not null" json:"-"`
|
||||
PublicKey string `gorm:"size:512;not null" json:"-"`
|
||||
AttestationType string `gorm:"size:64" json:"-"`
|
||||
AAGUID string `gorm:"size:64" json:"-"`
|
||||
SignCount uint64 `json:"-"`
|
||||
DeviceType string `gorm:"size:255" json:"device_type,omitempty"`
|
||||
LastUsedAt int64 `json:"last_used_at,omitempty"`
|
||||
BackupEligible bool `json:"-"`
|
||||
BackupState bool `json:"-"`
|
||||
Transport string `gorm:"size:32" json:"-"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// SystemConfig 系统配置
|
||||
type SystemConfig struct {
|
||||
Key string `gorm:"primaryKey;size:64" json:"key"`
|
||||
Value string `gorm:"type:jsonb;not null" json:"value"`
|
||||
}
|
||||
|
||||
// AllModels 返回所有需要迁移的模型
|
||||
func AllModels() []any {
|
||||
return []any{
|
||||
&User{},
|
||||
&APIKey{},
|
||||
&Channel{},
|
||||
&Model{},
|
||||
&ChannelModelBinding{},
|
||||
&UsageLog{},
|
||||
&UsageDaily{},
|
||||
&Passkey{},
|
||||
&SystemConfig{},
|
||||
}
|
||||
}
|
||||
|
||||
// HashAPIKey hashes an API key using SHA-256
|
||||
func HashAPIKey(key string) string {
|
||||
h := sha256.Sum256([]byte(key))
|
||||
return hex.EncodeToString(h[:])
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package usage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"opencatd-open/internal/dao"
|
||||
"opencatd-open/internal/store"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Event represents a usage event to be recorded
|
||||
type Event struct {
|
||||
UserID uint64
|
||||
ModelName string
|
||||
ChannelID uint64
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
CacheReadTokens int
|
||||
Cost float64
|
||||
IsError bool
|
||||
IsCanceled bool
|
||||
RequestID string
|
||||
}
|
||||
|
||||
// Recorder handles async usage recording
|
||||
type Recorder struct {
|
||||
usageDAO *dao.UsageDAO
|
||||
dailyDAO *dao.DailyUsageDAO
|
||||
ch chan Event
|
||||
batchSize int
|
||||
flushInterval time.Duration
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewRecorder creates a new usage recorder
|
||||
func NewRecorder(usageDAO *dao.UsageDAO, dailyDAO *dao.DailyUsageDAO) *Recorder {
|
||||
return &Recorder{
|
||||
usageDAO: usageDAO,
|
||||
dailyDAO: dailyDAO,
|
||||
ch: make(chan Event, 10000),
|
||||
batchSize: 100,
|
||||
flushInterval: 5 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// Start starts the recorder's background workers
|
||||
func (r *Recorder) Start(ctx context.Context) {
|
||||
r.wg.Add(1)
|
||||
go r.processLoop(ctx)
|
||||
}
|
||||
|
||||
// Stop gracefully stops the recorder
|
||||
func (r *Recorder) Stop() {
|
||||
close(r.ch)
|
||||
r.wg.Wait()
|
||||
}
|
||||
|
||||
// Record queues a usage event for async recording
|
||||
func (r *Recorder) Record(event Event) {
|
||||
select {
|
||||
case r.ch <- event:
|
||||
default:
|
||||
log.Printf("Usage channel full, dropping event for user %d model %s", event.UserID, event.ModelName)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) processLoop(ctx context.Context) {
|
||||
defer r.wg.Done()
|
||||
|
||||
batch := make([]Event, 0, r.batchSize)
|
||||
ticker := time.NewTicker(r.flushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
if len(batch) > 0 {
|
||||
r.flush(batch)
|
||||
}
|
||||
return
|
||||
case event, ok := <-r.ch:
|
||||
if !ok {
|
||||
if len(batch) > 0 {
|
||||
r.flush(batch)
|
||||
}
|
||||
return
|
||||
}
|
||||
batch = append(batch, event)
|
||||
if len(batch) >= r.batchSize {
|
||||
r.flush(batch)
|
||||
batch = make([]Event, 0, r.batchSize)
|
||||
}
|
||||
case <-ticker.C:
|
||||
if len(batch) > 0 {
|
||||
r.flush(batch)
|
||||
batch = make([]Event, 0, r.batchSize)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) flush(events []Event) {
|
||||
if len(events) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Batch create usage logs
|
||||
logs := make([]*store.UsageLog, 0, len(events))
|
||||
|
||||
for _, e := range events {
|
||||
status := store.UsageStatusSuccess
|
||||
if e.IsError {
|
||||
status = store.UsageStatusError
|
||||
}
|
||||
if e.IsCanceled {
|
||||
status = store.UsageStatusCanceled
|
||||
}
|
||||
|
||||
log := &store.UsageLog{
|
||||
UserID: e.UserID,
|
||||
ModelName: e.ModelName,
|
||||
ChannelID: e.ChannelID,
|
||||
InputTokens: int64(e.PromptTokens),
|
||||
OutputTokens: int64(e.CompletionTokens),
|
||||
CacheReadTokens: int64(e.CacheReadTokens),
|
||||
Cost: e.Cost,
|
||||
Status: status,
|
||||
RequestID: e.RequestID,
|
||||
}
|
||||
logs = append(logs, log)
|
||||
}
|
||||
|
||||
// Write to database
|
||||
if err := r.usageDAO.BatchCreate(context.Background(), logs); err != nil {
|
||||
log.Printf("Failed to batch create usage logs: %v", err)
|
||||
}
|
||||
|
||||
log.Printf("Flushed %d usage logs", len(logs))
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package utils
|
||||
|
||||
import "strings"
|
||||
|
||||
func StringToBool(strSlice []string) []bool {
|
||||
boolSlice := make([]bool, len(strSlice))
|
||||
for i, str := range strSlice {
|
||||
str = strings.ToLower(str)
|
||||
if str == "true" {
|
||||
boolSlice[i] = true
|
||||
} else if str == "false" {
|
||||
boolSlice[i] = false
|
||||
}
|
||||
}
|
||||
return boolSlice
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
package utils
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func MergeJSONObjects(dst, src map[string]interface{}) map[string]interface{} {
|
||||
|
||||
result := make(map[string]interface{})
|
||||
for k, v := range dst {
|
||||
result[k] = v
|
||||
}
|
||||
|
||||
for key, value2 := range src {
|
||||
value1, exists := result[key]
|
||||
|
||||
if exists {
|
||||
map1Val, map1IsMap := value1.(map[string]interface{})
|
||||
map2Val, map2IsMap := value2.(map[string]interface{})
|
||||
|
||||
if map1IsMap && map2IsMap {
|
||||
result[key] = MergeJSONObjects(map1Val, map2Val)
|
||||
} else {
|
||||
// 覆盖第一个map中的值
|
||||
result[key] = value2
|
||||
}
|
||||
} else {
|
||||
// 添加新的键值对
|
||||
result[key] = value2
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func StructToMap(in interface{}) (map[string]interface{}, error) {
|
||||
out := make(map[string]interface{})
|
||||
|
||||
v := reflect.ValueOf(in)
|
||||
// If it's a pointer, dereference it
|
||||
if v.Kind() == reflect.Ptr {
|
||||
v = v.Elem()
|
||||
}
|
||||
|
||||
// Check if it's a struct
|
||||
if v.Kind() != reflect.Struct {
|
||||
return nil, fmt.Errorf("StructToMap only accepts structs or pointers to structs; got %T", v.Interface())
|
||||
}
|
||||
|
||||
t := v.Type() // Get the type of the struct
|
||||
for i := 0; i < v.NumField(); i++ {
|
||||
// Get the field Value and Type
|
||||
fieldV := v.Field(i)
|
||||
fieldT := t.Field(i)
|
||||
|
||||
// Skip unexported fields
|
||||
if !fieldT.IsExported() {
|
||||
continue
|
||||
}
|
||||
|
||||
// --- Handle JSON Tag ---
|
||||
tag := fieldT.Tag.Get("json")
|
||||
key := fieldT.Name // Default key is the field name
|
||||
omitempty := false
|
||||
|
||||
if tag != "" {
|
||||
parts := strings.Split(tag, ",")
|
||||
tagName := parts[0]
|
||||
|
||||
if tagName == "-" {
|
||||
// Skip fields tagged with "-"
|
||||
continue
|
||||
}
|
||||
if tagName != "" {
|
||||
key = tagName // Use tag name as key
|
||||
}
|
||||
|
||||
// Check for omitempty option
|
||||
for _, part := range parts[1:] {
|
||||
if part == "omitempty" {
|
||||
omitempty = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Handle omitempty ---
|
||||
val := fieldV.Interface()
|
||||
if omitempty && fieldV.IsZero() {
|
||||
continue // Skip zero-value fields if omitempty is set
|
||||
}
|
||||
|
||||
// --- Handle Nested Structs/Pointers to Structs (Recursion) ---
|
||||
// Check for pointer first
|
||||
if fieldV.Kind() == reflect.Ptr {
|
||||
// If pointer is nil and omitempty is set, it was already skipped
|
||||
// If pointer is nil and omitempty is not set, add nil to map
|
||||
if fieldV.IsNil() {
|
||||
// Only add nil if omitempty is not set (already handled above)
|
||||
if !omitempty {
|
||||
out[key] = nil
|
||||
}
|
||||
continue // Move to next field
|
||||
}
|
||||
// If it points to a struct, dereference and recurse
|
||||
if fieldV.Elem().Kind() == reflect.Struct {
|
||||
nestedMap, err := StructToMap(fieldV.Interface()) // Pass the pointer
|
||||
if err != nil {
|
||||
// Decide how to handle nested errors, e.g., log or return
|
||||
fmt.Printf("Warning: could not convert nested struct pointer %s: %v\n", fieldT.Name, err)
|
||||
out[key] = val // Store original value on error? Or skip?
|
||||
} else {
|
||||
out[key] = nestedMap
|
||||
}
|
||||
continue // Move to next field after handling pointer
|
||||
}
|
||||
// If pointer to non-struct, just get the interface value (handled below)
|
||||
val = fieldV.Interface() // Use the actual pointer value
|
||||
|
||||
} else if fieldV.Kind() == reflect.Struct {
|
||||
// If it's a struct (not a pointer), recurse
|
||||
nestedMap, err := StructToMap(fieldV.Interface()) // Pass the struct value
|
||||
if err != nil {
|
||||
fmt.Printf("Warning: could not convert nested struct %s: %v\n", fieldT.Name, err)
|
||||
out[key] = val // Store original value on error? Or skip?
|
||||
} else {
|
||||
out[key] = nestedMap
|
||||
}
|
||||
continue // Move to next field after handling struct
|
||||
}
|
||||
|
||||
// Assign the value (primitive, slice, map, non-struct pointer, etc.)
|
||||
out[key] = val
|
||||
}
|
||||
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
package utils
|
||||
|
||||
import (
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
func HashPassword(password string) (string, error) {
|
||||
bytes, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
|
||||
return string(bytes), err
|
||||
}
|
||||
|
||||
func CheckPassword(hash, password string) bool {
|
||||
err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(password))
|
||||
return err == nil
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
package utils
|
||||
|
||||
func ToPtr[T any](v T) *T {
|
||||
return &v
|
||||
}
|
||||
|
||||
func UpdatePtrField[T any](target *T, value *T) {
|
||||
if value != nil {
|
||||
*target = *value
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user