Files
rag/model/entity/model.go
2026-06-10 16:36:46 +08:00

120 lines
3.5 KiB
Go

package entity
import (
"rag/consts/model"
"gitea.redpowerfuture.com/red-future/common/beans"
)
type modelCol struct {
beans.SQLBaseCol
DatasetId string
ModelType string
ModelName string
ModelDesc string
ConfigType string
ConfigContent string
}
var ModelCol = modelCol{
SQLBaseCol: beans.DefSQLBaseCol,
DatasetId: "dataset_id",
ModelType: "model_type",
ModelName: "model_name",
ModelDesc: "model_desc",
ConfigType: "config_type",
ConfigContent: "config_content",
}
type Model struct {
beans.SQLBaseDO `orm:",inline"`
DatasetId int64 `orm:"dataset_id" json:"datasetId" dc:"数据集ID"`
ModelType model.ModelType `orm:"model_type" json:"modelType" dc:"模型类型"` // 向量/对话
ModelName string `orm:"model_name" json:"modelName" dc:"模型名称"`
ModelDesc string `orm:"model_desc" json:"modelDesc" dc:"模型描述"`
ConfigType model.ModelConfigType `orm:"config_type" json:"configType" dc:"配置类型"` // ark/ollama等
ConfigContent map[string]interface{} `orm:"config_content" json:"configContent" dc:"配置详情"` // 存JSON
}
// -------------------------- 通用配置结构体(抽离重复字段)--------------------------
// OllamaConfig 通用配置(向量/对话完全一致)
type OllamaConfig struct {
BaseURL string `json:"base_url"`
Model string `json:"model"`
}
// OpenAIConfig 通用配置
type OpenAIConfig struct {
APIKey string `json:"api_key"`
Model string `json:"model"`
ByAzure bool `json:"by_azure"`
BaseURL string `json:"base_url"`
APIVersion string `json:"api_version"`
}
// QianfanConfig 千帆通用配置
type QianfanConfig struct {
AccessKey string `json:"access_key"`
SecretKey string `json:"secret_key"`
Model string `json:"model"`
}
// ArkConfig 通用配置
type ArkConfig struct {
APIKey string `json:"api_key"`
Model string `json:"model"`
}
// -------------------------- 向量模型配置 --------------------------
type VectorModelConfigOllama = OllamaConfig // 直接复用
type VectorModelConfigOpenAI = OpenAIConfig // 直接复用
type VectorModelConfigQianfan = QianfanConfig // 直接复用
type VectorModelConfigArk struct {
ArkConfig
APIType string `json:"api_type"`
}
type VectorModelConfigTencentCloud struct {
SecretID string `json:"secret_id"`
SecretKey string `json:"secret_key"`
Region string `json:"region"`
}
type VectorModelConfigDashScope struct {
APIKey string `json:"api_key"`
Model string `json:"model"`
}
// -------------------------- 对话模型配置 --------------------------
type ChatModelConfigArk = ArkConfig // 直接复用
type ChatModelConfigArkBot = ArkConfig // 直接复用
type ChatModelConfigOllama = OllamaConfig // 直接复用
type ChatModelConfigOpenAI = OpenAIConfig // 直接复用
type ChatModelConfigQianfan = QianfanConfig // 直接复用
type ChatModelConfigClaude struct {
ByBedrock bool `json:"by_bedrock"`
AccessKey string `json:"access_key"`
SecretAccessKey string `json:"secret_access_key"`
Region string `json:"region"`
APIKey string `json:"api_key"`
Model string `json:"model"`
BaseURL string `json:"base_url"`
}
type ChatModelConfigDeepSeek struct {
APIKey string `json:"api_key"`
Model string `json:"model"`
BaseURL string `json:"base_url"`
}
type ChatModelConfigQwen struct {
APIKey string `json:"api_key"`
Model string `json:"model"`
BaseURL string `json:"base_url"`
}