feat: 支持多租户多模型对话及文档去重优化
This commit is contained in:
243
common/eino/chat.go
Normal file
243
common/eino/chat.go
Normal file
@@ -0,0 +1,243 @@
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user