package types import ( "database/sql/driver" "encoding/json" "log" "time" "github.com/Tencent/WeKnora/internal/utils" "github.com/google/uuid" "gorm.io/gorm" ) // ModelType represents the type of AI model type ModelType string const ( ModelTypeEmbedding ModelType = "Embedding" // Embedding model ModelTypeRerank ModelType = "Rerank" // Rerank model ModelTypeKnowledgeQA ModelType = "KnowledgeQA" // KnowledgeQA model ModelTypeVLLM ModelType = "VLLM" // VLLM model ModelTypeASR ModelType = "ASR" // ASR (Automatic Speech Recognition) model ) // ModelStatus represents the status of the model type ModelStatus string const ( ModelStatusActive ModelStatus = "active" // Model is active ModelStatusDownloading ModelStatus = "downloading" // Model is downloading ModelStatusDownloadFailed ModelStatus = "download_failed" // Model download failed ) // ModelUsageBinding is a stable API identifier for one configuration field // that references a model. Clients localize these values; repository column // names and JSON paths must not leak into the response contract. type ModelUsageBinding string // Model usage binding constants identify supported model-reference roles. const ( ModelUsageBindingEmbeddingModel ModelUsageBinding = "embedding_model" ModelUsageBindingSummaryModel ModelUsageBinding = "summary_model" ModelUsageBindingImageProcessingModel ModelUsageBinding = "image_processing_model" ModelUsageBindingVLMModel ModelUsageBinding = "vlm_model" ModelUsageBindingASRModel ModelUsageBinding = "asr_model" ModelUsageBindingWikiSynthesisModel ModelUsageBinding = "wiki_synthesis_model" ModelUsageBindingChatModel ModelUsageBinding = "chat_model" ModelUsageBindingRerankModel ModelUsageBinding = "rerank_model" ModelUsageBindingQueryUnderstandModel ModelUsageBinding = "query_understand_model" ModelUsageBindingFollowUpModel ModelUsageBinding = "follow_up_model" ModelUsageBindingExtractModel ModelUsageBinding = "extract_model" ) // ModelUsageResource is the minimal, non-sensitive projection returned when a // knowledge base or agent prevents model deletion. type ModelUsageResource struct { ID string `json:"id"` Name string `json:"name"` Bindings []ModelUsageBinding `json:"bindings"` } // ModelUsageMemory describes workspace-level long-term-memory model pins. type ModelUsageMemory struct { Bindings []ModelUsageBinding `json:"bindings"` } // ModelUsageListLimit caps how many knowledge bases or agents are listed in a // model-in-use error. Totals are returned separately so the cap cannot // under-count the references that actually block deletion. const ModelUsageListLimit = 50 // ModelUsageDetails explains every active tenant-scoped reference that blocks // model deletion. KnowledgeBases and Agents may be truncated to // ModelUsageListLimit entries; KnowledgeBaseTotal and AgentTotal are the // untruncated counts used by the deletion guard and UI headings. type ModelUsageDetails struct { KnowledgeBases []ModelUsageResource `json:"knowledge_bases"` Agents []ModelUsageResource `json:"agents"` LongTermMemory ModelUsageMemory `json:"long_term_memory"` KnowledgeBaseTotal int64 `json:"knowledge_base_total"` AgentTotal int64 `json:"agent_total"` } // InUse reports whether at least one active object references the model. func (d ModelUsageDetails) InUse() bool { return d.KnowledgeBaseTotal > 0 || d.AgentTotal > 0 || len(d.KnowledgeBases) > 0 || len(d.Agents) > 0 || len(d.LongTermMemory.Bindings) > 0 } // ModelSource represents the source of the model type ModelSource string const ( ModelSourceLocal ModelSource = "local" // Local model ModelSourceRemote ModelSource = "remote" // Remote model ModelSourceAliyun ModelSource = "aliyun" // Aliyun DashScope model ModelSourceZhipu ModelSource = "zhipu" // Zhipu model ModelSourceVolcengine ModelSource = "volcengine" // Volcengine model ModelSourceDeepseek ModelSource = "deepseek" // Deepseek model ModelSourceHunyuan ModelSource = "hunyuan" // Hunyuan model ModelSourceMinimax ModelSource = "minimax" // Minimax mode ModelSourceOpenAI ModelSource = "openai" // OpenAI model ModelSourceGemini ModelSource = "gemini" // Gemini model ModelSourceMimo ModelSource = "mimo" // Mimo model ModelSourceSiliconFlow ModelSource = "siliconflow" // SiliconFlow model ModelSourceJina ModelSource = "jina" // Jina AI model ModelSourceOpenRouter ModelSource = "openrouter" // OpenRouter model ModelSourceLiteLLM ModelSource = "litellm" // LiteLLM proxy model ModelSourceRequesty ModelSource = "requesty" // Requesty model ModelSourceNvidia ModelSource = "nvidia" // NVIDIA model ModelSourceNovita ModelSource = "novita" // Novita AI model ModelSourceAzureOpenAI ModelSource = "azure_openai" // Azure OpenAI model ) // EmbeddingParameters represents the embedding parameters for a model type EmbeddingParameters struct { Dimension int `yaml:"dimension" json:"dimension"` TruncatePromptTokens int `yaml:"truncate_prompt_tokens" json:"truncate_prompt_tokens"` SupportsDimensionOverride bool `yaml:"supports_dimension_override" json:"supports_dimension_override"` } type ModelParameters struct { BaseURL string `yaml:"base_url" json:"base_url"` APIKey string `yaml:"api_key" json:"api_key"` InterfaceType string `yaml:"interface_type" json:"interface_type"` EmbeddingParameters EmbeddingParameters `yaml:"embedding_parameters" json:"embedding_parameters"` ParameterSize string `yaml:"parameter_size" json:"parameter_size"` // Ollama model parameter size (e.g., "7B", "13B", "70B") Provider string `yaml:"provider" json:"provider"` // Provider identifier: openai, aliyun, zhipu, generic ExtraConfig map[string]string `yaml:"extra_config" json:"extra_config"` // Provider-specific configuration // CustomHeaders 允许在调用远程模型 API 时附加自定义 HTTP 请求头, // 用途类似 Python OpenAI SDK 的 extra_headers 参数, // 常见场景包括透传企业网关鉴权信息、追踪 ID、路由标识等。 // 保留字段(Authorization、api-key、Content-Type、Accept 等)会在运行期被忽略以避免破坏签名/鉴权流程。 CustomHeaders map[string]string `yaml:"custom_headers,omitempty" json:"custom_headers,omitempty"` SupportsVision bool `yaml:"supports_vision" json:"supports_vision"` // Whether the model accepts image/multimodal input // ContextWindow is the model's total context window in tokens and // MaxOutputTokens the most it emits in one response. Both are provider // facts the agent cannot discover but has to act on: the context window is // what decides when conversation history gets compacted, and assuming a // window larger than the real one means compaction never fires and the // provider rejects the request mid-conversation instead. 0 means unknown, // which falls back to DefaultMaxContextTokens. ContextWindow int `yaml:"context_window,omitempty" json:"context_window,omitempty"` MaxOutputTokens int `yaml:"max_output_tokens,omitempty" json:"max_output_tokens,omitempty"` // MaxConcurrency caps concurrent in-flight BACKGROUND (ingestion / // enrichment) calls to THIS specific model, keyed by model ID and shared // across all replicas. 0 (the default) means "fall back to the // process-wide model.max_concurrency". Interactive user-facing calls are // never gated. Only chat / vlm / embedding honour this (see limiter.Gate). MaxConcurrency int `yaml:"max_concurrency,omitempty" json:"max_concurrency,omitempty"` // WeKnoraCloud 厂商专用凭证 AppID string `yaml:"app_id,omitempty" json:"app_id,omitempty"` AppSecret string `yaml:"app_secret,omitempty" json:"app_secret,omitempty"` // AES-256 加密存储,实际承载上游 API Key } // Per-response redaction for Model now lives in dto.NewModelResponse. The // previous RedactSensitiveData method has been removed because handlers must // always serialize through the DTO, where the secret fields don't even // exist; a runtime mutator on the entity is both redundant and a footgun // (mutates an entity that other code may still be using). // ModelIDMaxLen is the upper bound on `models.id`. Matches the actual // schema width on both PostgreSQL (varchar(64) in migrations/versioned/ // 000000_init.up.sql) and SQLite (varchar(64) in migrations/sqlite/ // 000000_init.up.sql). Loaders that accept user-provided ids (e.g. the // built-in models YAML loader) must reject anything longer to avoid a // "value too long for type" failure at INSERT time. const ModelIDMaxLen = 64 // DefaultBuiltinModelTenantID is the tenant id that built-in models are // assigned to when YAML does not specify one. Kept in sync with the seed // value of tenants_id_seq in migrations/versioned/000000_init.up.sql // (and the equivalent SQLite init); changing one without the other will // break visibility of built-in models for the default tenant. const DefaultBuiltinModelTenantID uint64 = 10000 // Model represents the AI model type Model struct { // Unique identifier of the model. The actual DB schema width is // varchar(64) on both PostgreSQL and SQLite (see ModelIDMaxLen); // GORM's struct tag is documented to match so AutoMigrate paths // produce the same shape. ID string `yaml:"id" json:"id" gorm:"type:varchar(64);primaryKey"` // Workspace ID TenantID uint64 `yaml:"tenant_id" json:"tenant_id"` // Name of the model Name string `yaml:"name" json:"name"` // Optional user-facing display name. Runtime calls still use Name. DisplayName string `yaml:"display_name" json:"display_name" gorm:"type:varchar(255);default:''"` // Type of the model Type ModelType `yaml:"type" json:"type"` // Source of the model Source ModelSource `yaml:"source" json:"source"` // Description of the model Description string `yaml:"description" json:"description"` // Model parameters in JSON format Parameters ModelParameters `yaml:"parameters" json:"parameters" gorm:"type:json"` // Whether the model is the default model IsDefault bool `yaml:"is_default" json:"is_default"` // Whether the model is a builtin model (visible to all tenants) IsBuiltin bool `yaml:"is_builtin" json:"is_builtin" gorm:"default:false"` // ManagedBy identifies which subsystem owns this row's lifecycle. // Empty / "" = manually created (UI / API / hand-written SQL); the YAML // builtin-models loader leaves these untouched. // "yaml" = declared in config/builtin_models.yaml; on every startup the // loader UPSERTs the YAML set and soft-deletes YAML-managed rows whose // id is no longer present in the file. Future origins (e.g. "helm", // "operator") can claim their own slice without interfering. ManagedBy string `yaml:"managed_by" json:"managed_by,omitempty" gorm:"type:varchar(32);default:''"` // Model status, default: active, possible: downloading, download_failed Status ModelStatus `yaml:"status" json:"status"` // Creation time of the model CreatedAt time.Time `yaml:"created_at" json:"created_at"` // Last updated time of the model UpdatedAt time.Time `yaml:"updated_at" json:"updated_at"` // Deletion time of the model DeletedAt gorm.DeletedAt `yaml:"deleted_at" json:"deleted_at" gorm:"index"` } // Value implements the driver.Valuer interface, used to convert ModelParameters to database value. // Encrypts APIKey and AppSecret before persisting to database (value receiver = no memory pollution). func (c ModelParameters) Value() (driver.Value, error) { if key := utils.GetAESKey(); key != nil { if c.APIKey != "" { if encrypted, err := utils.EncryptAESGCM(c.APIKey, key); err == nil { c.APIKey = encrypted } } if c.AppSecret != "" { if encrypted, err := utils.EncryptAESGCM(c.AppSecret, key); err == nil { c.AppSecret = encrypted } } } return json.Marshal(c) } // Scan implements the sql.Scanner interface, used to convert database value to ModelParameters. // Decrypts APIKey and AppSecret after loading from database; legacy plaintext is returned as-is. func (c *ModelParameters) Scan(value interface{}) error { if value == nil { return nil } b, ok := value.([]byte) if !ok { return nil } if err := json.Unmarshal(b, c); err != nil { return err } // Lenient decrypt: a row with broken ciphertext (key rotated/removed) // must still load — otherwise a single failure breaks ListModels and // the user can't even see which model needs re-credentialing. if plain, ok := utils.DecryptStoredSecretLenient(c.APIKey); ok { c.APIKey = plain } else { log.Printf("[crypto] model parameters api_key: decrypt failed (SYSTEM_AES_KEY missing/rotated?), treating as unconfigured") c.APIKey = "" } if plain, ok := utils.DecryptStoredSecretLenient(c.AppSecret); ok { c.AppSecret = plain } else { log.Printf("[crypto] model parameters app_secret: decrypt failed (SYSTEM_AES_KEY missing/rotated?), treating as unconfigured") c.AppSecret = "" } return nil } // BeforeCreate is a GORM hook that runs before creating a new model record. // Generates a UUID only when the caller has not supplied an ID — preserves // stable IDs declared in built-in model YAML config while keeping the // existing UUID behaviour for API-driven model creation. func (m *Model) BeforeCreate(tx *gorm.DB) (err error) { if m.ID == "" { m.ID = uuid.New().String() } return nil }