103 lines
2.5 KiB
Go
103 lines
2.5 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/cloudwego/eino-ext/components/model/qwen"
|
|
"github.com/cloudwego/eino/components/model"
|
|
"github.com/gogf/gf/v2/util/gconv"
|
|
)
|
|
|
|
// ModelConfig 模型配置 — 所有字段必须显式提供，无硬编码默认值
|
|
type ModelConfig struct {
|
|
ModelName string // 对话模型名
|
|
APIKey string // API密钥
|
|
BaseURL string // API地址
|
|
MaxTokens int // 最大Token数
|
|
Temperature float32 // 温度参数
|
|
}
|
|
|
|
// context keys
|
|
type ctxKey string
|
|
|
|
const (
|
|
ctxKeyAPIKey ctxKey = "api_key"
|
|
ctxKeyBaseURL ctxKey = "base_url"
|
|
ctxKeyDramaId ctxKey = "drama_id"
|
|
)
|
|
|
|
// WithDramaId 将短剧ID注入 context，供工具函数读取场景图片
|
|
func WithDramaId(ctx context.Context, dramaId int64) context.Context {
|
|
return context.WithValue(ctx, ctxKeyDramaId, dramaId)
|
|
}
|
|
|
|
// GetDramaId 从 context 获取短剧ID
|
|
func GetDramaId(ctx context.Context) int64 {
|
|
if v, ok := ctx.Value(ctxKeyDramaId).(int64); ok {
|
|
return v
|
|
}
|
|
return 0
|
|
}
|
|
|
|
// WithModelConfig 将模型配置注入 context，供工具函数读取
|
|
func WithModelConfig(ctx context.Context, cfg *ModelConfig) context.Context {
|
|
ctx = context.WithValue(ctx, ctxKeyAPIKey, cfg.APIKey)
|
|
ctx = context.WithValue(ctx, ctxKeyBaseURL, cfg.BaseURL)
|
|
return ctx
|
|
}
|
|
|
|
// GetAPIKey 从 context 获取 API key
|
|
func GetAPIKey(ctx context.Context) string {
|
|
if v, ok := ctx.Value(ctxKeyAPIKey).(string); ok && v != "" {
|
|
return v
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// GetBaseURL 从 context 获取 API 地址
|
|
func GetBaseURL(ctx context.Context) string {
|
|
if v, ok := ctx.Value(ctxKeyBaseURL).(string); ok && v != "" {
|
|
return v
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// NewChatModel 根据配置初始化聊天模型
|
|
func NewChatModel(ctx context.Context, cfg *ModelConfig) (cm model.ChatModel, err error) {
|
|
if cfg == nil {
|
|
return nil, fmt.Errorf("模型配置不能为空")
|
|
}
|
|
if cfg.APIKey == "" {
|
|
return nil, fmt.Errorf("APIKey 未配置")
|
|
}
|
|
if cfg.ModelName == "" {
|
|
return nil, fmt.Errorf("模型名称未配置")
|
|
}
|
|
if cfg.BaseURL == "" {
|
|
return nil, fmt.Errorf("API 地址未配置")
|
|
}
|
|
|
|
maxTokens := cfg.MaxTokens
|
|
if maxTokens <= 0 {
|
|
maxTokens = 4096
|
|
}
|
|
temperature := cfg.Temperature
|
|
if temperature <= 0 {
|
|
temperature = 0.8
|
|
}
|
|
|
|
config := &qwen.ChatModelConfig{
|
|
APIKey: cfg.APIKey,
|
|
Model: cfg.ModelName,
|
|
BaseURL: cfg.BaseURL,
|
|
MaxTokens: gconv.PtrInt(maxTokens),
|
|
Temperature: gconv.PtrFloat32(temperature),
|
|
}
|
|
cm, err = qwen.NewChatModel(ctx, config)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("创建模型失败: %w", err)
|
|
}
|
|
return cm, nil
|
|
}
|