244 lines
7.6 KiB
Go
244 lines
7.6 KiB
Go
|
package eino
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"fmt"
|
|||
|
|
"rag/consts/model"
|
|||
|
|
"rag/dao"
|
|||
|
|
"rag/model/dto"
|
|||
|
|
"rag/model/entity"
|
|||
|
|
|
|||
|
|
"gitea.com/red-future/common/jaeger"
|
|||
|
|
"gitea.com/red-future/common/utils"
|
|||
|
|
"github.com/cloudwego/eino-ext/components/model/ark"
|
|||
|
|
"github.com/cloudwego/eino-ext/components/model/arkbot"
|
|||
|
|
"github.com/cloudwego/eino-ext/components/model/claude"
|
|||
|
|
"github.com/cloudwego/eino-ext/components/model/deepseek"
|
|||
|
|
"github.com/cloudwego/eino-ext/components/model/ollama"
|
|||
|
|
"github.com/cloudwego/eino-ext/components/model/openai"
|
|||
|
|
"github.com/cloudwego/eino-ext/components/model/qianfan"
|
|||
|
|
"github.com/cloudwego/eino-ext/components/model/qwen"
|
|||
|
|
modelChat "github.com/cloudwego/eino/components/model"
|
|||
|
|
"github.com/gogf/gf/v2/frame/g"
|
|||
|
|
"github.com/gogf/gf/v2/util/gconv"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
type ChatModelSet struct {
|
|||
|
|
Ark *ark.ChatModel
|
|||
|
|
ArkBot *arkbot.ChatModel
|
|||
|
|
Claude *claude.ChatModel
|
|||
|
|
DeepSeek *deepseek.ChatModel
|
|||
|
|
Ollama *ollama.ChatModel
|
|||
|
|
OpenAI *openai.ChatModel
|
|||
|
|
Qianfan *qianfan.ChatModel
|
|||
|
|
Qwen *qwen.ChatModel
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 全局租户容器：key=tenantId，value=该租户的对话模型
|
|||
|
|
var tenantChatModels = make(map[uint64]*ChatModelSet)
|
|||
|
|
|
|||
|
|
func init() {
|
|||
|
|
ctx := context.Background()
|
|||
|
|
ctx, span := jaeger.NewSpan(ctx, "InitAllChat")
|
|||
|
|
defer span.End()
|
|||
|
|
InitAllChat(ctx)
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ===================== 1. 服务启动时：初始化所有租户对话模型 =====================
|
|||
|
|
func InitAllChat(ctx context.Context) {
|
|||
|
|
list, err := dao.Model.GetNoTenantId(ctx, &dto.GetModelReq{
|
|||
|
|
ModelType: model.ModelTypeChat.Code(),
|
|||
|
|
})
|
|||
|
|
if err != nil {
|
|||
|
|
g.Log().Errorf(ctx, "获取所有租户对话模型失败: %v", err)
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
for _, l := range list {
|
|||
|
|
err = InitChat(ctx, l)
|
|||
|
|
if err != nil {
|
|||
|
|
g.Log().Errorf(ctx, "初始化租户[%v]的对话模型失败: %v", l.TenantId, err)
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func InitChat(ctx context.Context, modelDO *entity.Model) (err error) {
|
|||
|
|
set := &ChatModelSet{}
|
|||
|
|
switch *modelDO.ConfigType {
|
|||
|
|
case *model.ModelConfigTypeChatArk.Code():
|
|||
|
|
var cfg entity.ChatModelConfigArk
|
|||
|
|
if err = gconv.Struct(modelDO.ConfigContent, &cfg); err != nil {
|
|||
|
|
return fmt.Errorf("解析Ark配置失败: %v", err)
|
|||
|
|
}
|
|||
|
|
set.Ark, err = ark.NewChatModel(ctx, &ark.ChatModelConfig{
|
|||
|
|
APIKey: cfg.APIKey,
|
|||
|
|
Model: cfg.Model,
|
|||
|
|
Temperature: gconv.PtrFloat32(0.7),
|
|||
|
|
MaxTokens: gconv.PtrInt(1024),
|
|||
|
|
TopP: gconv.PtrFloat32(1.0),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
case *model.ModelConfigTypeChatArkBot.Code():
|
|||
|
|
var cfg entity.ChatModelConfigArkBot
|
|||
|
|
if err = gconv.Struct(modelDO.ConfigContent, &cfg); err != nil {
|
|||
|
|
return fmt.Errorf("解析ArkBot配置失败: %v", err)
|
|||
|
|
}
|
|||
|
|
set.ArkBot, err = arkbot.NewChatModel(ctx, &arkbot.Config{
|
|||
|
|
APIKey: cfg.APIKey,
|
|||
|
|
Model: cfg.Model,
|
|||
|
|
Temperature: gconv.PtrFloat32(0.7),
|
|||
|
|
MaxTokens: gconv.PtrInt(1024),
|
|||
|
|
TopP: gconv.PtrFloat32(1.0),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
case *model.ModelConfigTypeChatClaude.Code():
|
|||
|
|
var cfg entity.ChatModelConfigClaude
|
|||
|
|
if err = gconv.Struct(modelDO.ConfigContent, &cfg); err != nil {
|
|||
|
|
return fmt.Errorf("解析Claude配置失败: %v", err)
|
|||
|
|
}
|
|||
|
|
claudeCfg := claude.Config{
|
|||
|
|
APIKey: cfg.APIKey,
|
|||
|
|
BaseURL: gconv.PtrString(cfg.BaseURL),
|
|||
|
|
Model: cfg.Model,
|
|||
|
|
Temperature: gconv.PtrFloat32(0.7),
|
|||
|
|
MaxTokens: gconv.Int(1024),
|
|||
|
|
TopP: gconv.PtrFloat32(1.0),
|
|||
|
|
ByBedrock: cfg.ByBedrock,
|
|||
|
|
AccessKey: cfg.AccessKey,
|
|||
|
|
SecretAccessKey: cfg.SecretAccessKey,
|
|||
|
|
Region: cfg.Region,
|
|||
|
|
}
|
|||
|
|
set.Claude, err = claude.NewChatModel(ctx, &claudeCfg)
|
|||
|
|
|
|||
|
|
case *model.ModelConfigTypeChatDeepSeek.Code():
|
|||
|
|
var cfg entity.ChatModelConfigDeepSeek
|
|||
|
|
if err = gconv.Struct(modelDO.ConfigContent, &cfg); err != nil {
|
|||
|
|
return fmt.Errorf("解析DeepSeek配置失败: %v", err)
|
|||
|
|
}
|
|||
|
|
set.DeepSeek, err = deepseek.NewChatModel(ctx, &deepseek.ChatModelConfig{
|
|||
|
|
APIKey: cfg.APIKey,
|
|||
|
|
Model: cfg.Model,
|
|||
|
|
BaseURL: cfg.BaseURL,
|
|||
|
|
Temperature: gconv.Float32(0.7),
|
|||
|
|
MaxTokens: gconv.Int(1024),
|
|||
|
|
TopP: gconv.Float32(1.0),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
case *model.ModelConfigTypeChatOllama.Code():
|
|||
|
|
var cfg entity.ChatModelConfigOllama
|
|||
|
|
if err = gconv.Struct(modelDO.ConfigContent, &cfg); err != nil {
|
|||
|
|
return fmt.Errorf("解析Ollama配置失败: %v", err)
|
|||
|
|
}
|
|||
|
|
set.Ollama, err = ollama.NewChatModel(ctx, &ollama.ChatModelConfig{
|
|||
|
|
BaseURL: cfg.BaseURL,
|
|||
|
|
Model: cfg.Model,
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
case *model.ModelConfigTypeChatOpenAI.Code():
|
|||
|
|
var cfg entity.ChatModelConfigOpenAI
|
|||
|
|
if err = gconv.Struct(modelDO.ConfigContent, &cfg); err != nil {
|
|||
|
|
return fmt.Errorf("解析OpenAI配置失败: %v", err)
|
|||
|
|
}
|
|||
|
|
openAiCfg := openai.ChatModelConfig{
|
|||
|
|
APIKey: cfg.APIKey,
|
|||
|
|
Model: cfg.Model,
|
|||
|
|
ByAzure: cfg.ByAzure,
|
|||
|
|
BaseURL: cfg.BaseURL,
|
|||
|
|
APIVersion: cfg.APIVersion,
|
|||
|
|
Temperature: gconv.PtrFloat32(0.7),
|
|||
|
|
MaxCompletionTokens: gconv.PtrInt(1024),
|
|||
|
|
TopP: gconv.PtrFloat32(1.0),
|
|||
|
|
}
|
|||
|
|
set.OpenAI, err = openai.NewChatModel(ctx, &openAiCfg)
|
|||
|
|
|
|||
|
|
case *model.ModelConfigTypeChatQianfan.Code():
|
|||
|
|
var cfg entity.ChatModelConfigQianfan
|
|||
|
|
if err = gconv.Struct(modelDO.ConfigContent, &cfg); err != nil {
|
|||
|
|
return fmt.Errorf("解析千帆配置失败: %v", err)
|
|||
|
|
}
|
|||
|
|
qcfg := qianfan.GetQianfanSingletonConfig()
|
|||
|
|
qcfg.AccessKey = cfg.AccessKey
|
|||
|
|
qcfg.SecretKey = cfg.SecretKey
|
|||
|
|
set.Qianfan, err = qianfan.NewChatModel(ctx, &qianfan.ChatModelConfig{
|
|||
|
|
Model: cfg.Model,
|
|||
|
|
Temperature: gconv.PtrFloat32(0.7),
|
|||
|
|
MaxCompletionTokens: gconv.PtrInt(1024),
|
|||
|
|
TopP: gconv.PtrFloat32(1.0),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
case *model.ModelConfigTypeChatQwen.Code():
|
|||
|
|
var cfg entity.ChatModelConfigQwen
|
|||
|
|
if err = gconv.Struct(modelDO.ConfigContent, &cfg); err != nil {
|
|||
|
|
return fmt.Errorf("解析Qwen配置失败: %v", err)
|
|||
|
|
}
|
|||
|
|
set.Qwen, err = qwen.NewChatModel(ctx, &qwen.ChatModelConfig{
|
|||
|
|
APIKey: cfg.APIKey,
|
|||
|
|
Model: cfg.Model,
|
|||
|
|
BaseURL: cfg.BaseURL,
|
|||
|
|
Temperature: gconv.PtrFloat32(0.7),
|
|||
|
|
MaxTokens: gconv.PtrInt(1024),
|
|||
|
|
TopP: gconv.PtrFloat32(1.0),
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
default:
|
|||
|
|
return fmt.Errorf("不支持的对话模型类型: %v", *modelDO.ConfigType)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if err != nil {
|
|||
|
|
return fmt.Errorf("初始化对话模型失败: %v", err)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 无锁存入租户 map
|
|||
|
|
tenantChatModels[modelDO.TenantId] = set
|
|||
|
|
g.Log().Infof(ctx, "租户[%v]对话模型[%v]初始化成功", modelDO.TenantId, *modelDO.ConfigType)
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func GetTenantChatModel(tenantId uint64) (*ChatModelSet, error) {
|
|||
|
|
set := tenantChatModels[tenantId]
|
|||
|
|
if set == nil {
|
|||
|
|
return nil, fmt.Errorf("租户[%v]对话模型未初始化", tenantId)
|
|||
|
|
}
|
|||
|
|
return set, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func GetTenantChatModelByType(ctx context.Context, configType model.ModelConfigType) (modelChat.BaseChatModel, error) {
|
|||
|
|
userInfo, err := utils.GetUserInfo(ctx)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
set, err := GetTenantChatModel(userInfo.TenantId)
|
|||
|
|
if set == nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
switch *configType {
|
|||
|
|
case *model.ModelConfigTypeChatArk.Code():
|
|||
|
|
return set.Ark, nil
|
|||
|
|
case *model.ModelConfigTypeChatArkBot.Code():
|
|||
|
|
return set.ArkBot, nil
|
|||
|
|
case *model.ModelConfigTypeChatClaude.Code():
|
|||
|
|
return set.Claude, nil
|
|||
|
|
case *model.ModelConfigTypeChatDeepSeek.Code():
|
|||
|
|
return set.DeepSeek, nil
|
|||
|
|
case *model.ModelConfigTypeChatOllama.Code():
|
|||
|
|
return set.Ollama, nil
|
|||
|
|
case *model.ModelConfigTypeChatOpenAI.Code():
|
|||
|
|
return set.OpenAI, nil
|
|||
|
|
case *model.ModelConfigTypeChatQianfan.Code():
|
|||
|
|
return set.Qianfan, nil
|
|||
|
|
case *model.ModelConfigTypeChatQwen.Code():
|
|||
|
|
return set.Qwen, nil
|
|||
|
|
default:
|
|||
|
|
return nil, fmt.Errorf("不支持的对话模型类型: %v", configType)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func RefreshTenantChatModel(ctx context.Context, modelDO *entity.Model) error {
|
|||
|
|
delete(tenantChatModels, modelDO.TenantId)
|
|||
|
|
return InitChat(ctx, modelDO)
|
|||
|
|
}
|