上传初始版，功能未经任何测试，仅为基本框架，目前为完全不可用状态
This commit is contained in:
@@ -0,0 +1,132 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"d:\项目\dobaochet\backend\config"
|
||||
)
|
||||
|
||||
// AIService AI服务结构体
|
||||
type AIService struct {
|
||||
cfg *config.Config
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// NewAIService 创建AI服务实例
|
||||
func NewAIService(cfg *config.Config) *AIService {
|
||||
return &AIService{
|
||||
cfg: cfg,
|
||||
client: &http.Client{},
|
||||
}
|
||||
}
|
||||
|
||||
// AIChatRequest AI聊天请求结构体
|
||||
type AIChatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []Message `json:"messages"`
|
||||
MaxTokens int64 `json:"max_tokens"`
|
||||
Temperature float64 `json:"temperature"`
|
||||
}
|
||||
|
||||
// Message 消息结构体
|
||||
type Message struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// AIChatResponse AI聊天响应结构体
|
||||
type AIChatResponse 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"`
|
||||
}
|
||||
|
||||
// Choice 响应选项结构体
|
||||
type Choice struct {
|
||||
Index int `json:"index"`
|
||||
Message Message `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
|
||||
// Usage 用量结构体
|
||||
type Usage struct {
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
}
|
||||
|
||||
// Chat 调用AI模型聊天
|
||||
func (s *AIService) Chat(messages []Message) (*AIChatResponse, error) {
|
||||
// 构建请求
|
||||
reqBody := AIChatRequest{
|
||||
Model: s.cfg.AI.ModelName,
|
||||
Messages: messages,
|
||||
MaxTokens: s.cfg.AI.MaxTokens,
|
||||
Temperature: s.cfg.AI.Temperature,
|
||||
}
|
||||
|
||||
// 序列化请求体
|
||||
reqBytes, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("序列化请求失败: %w", err)
|
||||
}
|
||||
|
||||
// 创建HTTP请求
|
||||
req, err := http.NewRequest("POST", s.cfg.AI.APIURL, bytes.NewBuffer(reqBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建请求失败: %w", err)
|
||||
}
|
||||
|
||||
// 设置请求头
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if s.cfg.AI.APIKey != "" {
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", s.cfg.AI.APIKey))
|
||||
}
|
||||
|
||||
// 发送请求
|
||||
resp, err := s.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("发送请求失败: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// 读取响应
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取响应失败: %w", err)
|
||||
}
|
||||
|
||||
// 检查响应状态
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("AI API返回错误: %s, 响应: %s", resp.Status, string(respBody))
|
||||
}
|
||||
|
||||
// 解析响应
|
||||
var aiResp AIChatResponse
|
||||
if err := json.Unmarshal(respBody, &aiResp); err != nil {
|
||||
return nil, fmt.Errorf("解析响应失败: %w", err)
|
||||
}
|
||||
|
||||
return &aiResp, nil
|
||||
}
|
||||
|
||||
// CountTokens 估算令牌数量（简化实现，实际应使用更准确的方法）
|
||||
func (s *AIService) CountTokens(text string) int64 {
|
||||
// 简单估算：每个汉字算2个令牌，每个英文单词算1个令牌
|
||||
var count int64
|
||||
for _, r := range text {
|
||||
if r > 127 {
|
||||
count += 2
|
||||
} else {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count / 2 // 平均估算
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"d:\项目\dobaochet\backend\config"
|
||||
"d:\项目\dobaochet\backend\db"
|
||||
"d:\项目\dobaochet\backend\models"
|
||||
)
|
||||
|
||||
// QuotaService 配额管理服务
|
||||
type QuotaService struct {
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
// NewQuotaService 创建配额管理服务实例
|
||||
func NewQuotaService(cfg *config.Config) *QuotaService {
|
||||
return &QuotaService{cfg: cfg}
|
||||
}
|
||||
|
||||
// StartQuotaResetJob 启动配额重置定时任务
|
||||
func (s *QuotaService) StartQuotaResetJob() {
|
||||
// 立即执行一次重置检查
|
||||
s.ResetExpiredQuotas()
|
||||
|
||||
// 每天检查一次
|
||||
ticker := time.NewTicker(24 * time.Hour)
|
||||
go func() {
|
||||
for range ticker.C {
|
||||
s.ResetExpiredQuotas()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// ResetExpiredQuotas 重置过期的配额
|
||||
func (s *QuotaService) ResetExpiredQuotas() {
|
||||
log.Println("开始检查并重置过期配额")
|
||||
|
||||
// 查询所有需要重置配额的用户
|
||||
var users []models.User
|
||||
result := db.GetDB().Preload("Quota").Where("quota.reset_at <= ?", time.Now()).Find(&users)
|
||||
if result.Error != nil {
|
||||
log.Printf("查询需要重置配额的用户失败: %v", result.Error)
|
||||
return
|
||||
}
|
||||
|
||||
if len(users) == 0 {
|
||||
log.Println("没有需要重置配额的用户")
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("发现 %d 个用户需要重置配额", len(users))
|
||||
|
||||
// 批量重置配额
|
||||
tx := db.GetDB().Begin()
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
tx.Rollback()
|
||||
log.Printf("重置配额时发生 panic: %v", r)
|
||||
}
|
||||
}()
|
||||
|
||||
for _, user := range users {
|
||||
// 重置配额
|
||||
user.Quota.UsedTokens = 0
|
||||
user.Quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)
|
||||
|
||||
if result := tx.Save(&user.Quota); result.Error != nil {
|
||||
tx.Rollback()
|
||||
log.Printf("重置用户 %s 配额失败: %v", user.Username, result.Error)
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("已重置用户 %s 的配额", user.Username)
|
||||
}
|
||||
|
||||
if err := tx.Commit().Error; err != nil {
|
||||
log.Printf("提交配额重置事务失败: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("成功重置 %d 个用户的配额", len(users))
|
||||
}
|
||||
|
||||
// CheckAndUpdateQuota 检查并更新用户配额
|
||||
func (s *QuotaService) CheckAndUpdateQuota(userID uint, usedTokens int64) (bool, error) {
|
||||
// 查询用户配额
|
||||
var quota models.Quota
|
||||
result := db.GetDB().Where("user_id = ?", userID).First("a)
|
||||
if result.Error != nil {
|
||||
return false, result.Error
|
||||
}
|
||||
|
||||
// 检查配额是否过期
|
||||
if quota.ResetAt.Before(time.Now()) {
|
||||
// 重置配额
|
||||
quota.UsedTokens = 0
|
||||
quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)
|
||||
}
|
||||
|
||||
// 检查配额是否足够
|
||||
if quota.UsedTokens+usedTokens > quota.TotalTokens {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// 更新配额
|
||||
quota.UsedTokens += usedTokens
|
||||
if result := db.GetDB().Save("a); result.Error != nil {
|
||||
return false, result.Error
|
||||
}
|
||||
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// GetUserQuota 获取用户配额信息
|
||||
func (s *QuotaService) GetUserQuota(userID uint) (*models.Quota, error) {
|
||||
var quota models.Quota
|
||||
result := db.GetDB().Where("user_id = ?", userID).First("a)
|
||||
if result.Error != nil {
|
||||
return nil, result.Error
|
||||
}
|
||||
|
||||
// 检查并重置过期配额
|
||||
if quota.ResetAt.Before(time.Now()) {
|
||||
quota.UsedTokens = 0
|
||||
quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)
|
||||
db.GetDB().Save("a)
|
||||
}
|
||||
|
||||
return "a, nil
|
||||
}
|
||||
|
||||
// UpdateUserQuota 更新用户配额
|
||||
func (s *QuotaService) UpdateUserQuota(userID uint, totalTokens, tokenLimit int64) error {
|
||||
var quota models.Quota
|
||||
result := db.GetDB().Where("user_id = ?", userID).First("a)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// 更新配额
|
||||
quota.TotalTokens = totalTokens
|
||||
quota.TokenLimit = tokenLimit
|
||||
|
||||
return db.GetDB().Save("a).Error
|
||||
}
|
||||
|
||||
// ResetUserQuota 重置单个用户配额
|
||||
func (s *QuotaService) ResetUserQuota(userID uint) error {
|
||||
var quota models.Quota
|
||||
result := db.GetDB().Where("user_id = ?", userID).First("a)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
|
||||
// 重置配额
|
||||
quota.UsedTokens = 0
|
||||
quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)
|
||||
|
||||
return db.GetDB().Save("a).Error
|
||||
}
|
||||
Reference in New Issue
Block a user