diff --git a/docker/.env b/docker/.env index fa1b7693d0..e6bab62333 100644 --- a/docker/.env +++ b/docker/.env @@ -48,7 +48,7 @@ DEVICE=${DEVICE:-cpu} # This controls only the bundled metadata service started by Compose; DB_TYPE # independently selects the metadata backend used by RAGFlow. Keep `mysql` to # start the in-cluster MySQL service. For external GaussDB, set both -# DB_TYPE=gaussdb and METADATA_DB_PROFILE=gaussdb so MySQL stays disabled. +# DB_TYPE=gaussdb and METADATA_DB_PROFILE=gaussdb so MySQL stays disabled. # For vastbase serving as both doc engine and metadata database, set DB_TYPE=vastbase # and METADATA_DB_PROFILE=vastbase so MySQL stays disabled. METADATA_DB_PROFILE=${METADATA_DB_PROFILE:-mysql} diff --git a/docker/service_conf.yaml.template b/docker/service_conf.yaml.template index 06881e19af..55911ef09a 100644 --- a/docker/service_conf.yaml.template +++ b/docker/service_conf.yaml.template @@ -261,4 +261,4 @@ user_default_llm: # path: audit_log.log # empty/missing = binary default (server_main.log / admin_server.log / ingestion_server.log) # max_size: 100 # MB before rotation # max_backups: 10 # retained rotated files -# max_age: 30 # days \ No newline at end of file +# max_age: 30 # days diff --git a/internal/agent/component/categorize_decision_test.go b/internal/agent/component/categorize_decision_test.go index 6d24ceb7ee..50257af836 100644 --- a/internal/agent/component/categorize_decision_test.go +++ b/internal/agent/component/categorize_decision_test.go @@ -24,7 +24,7 @@ func setupCategorizeDecisionModel(t *testing.T, baseURL string) *gorm.DB { dao.GetModelProviderManager() db := setupComponentTestDB(t) for _, row := range []any{ - &entity.TenantModelProvider{ID: "decision-provider", TenantID: "tenant-1", ProviderName: "SystemOne"}, + &entity.TenantModelProvider{ID: "decision-provider", UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "SystemOne"}, &entity.TenantModelInstance{ID: "decision-instance", ProviderID: "decision-provider", InstanceName: "default", APIKey: "test-key", Status: "active", Extra: fmt.Sprintf(`{"base_url":%q}`, baseURL)}, &entity.TenantModel{ID: categorizeDecisionModelID, ProviderID: "decision-provider", InstanceID: "decision-instance", ModelName: "decision-model-preview", ModelType: int(entity.ModelTypeDecision), Status: "active"}, } { diff --git a/internal/agent/component/categorize_test.go b/internal/agent/component/categorize_test.go index 6fef38e967..a949f0b21c 100644 --- a/internal/agent/component/categorize_test.go +++ b/internal/agent/component/categorize_test.go @@ -379,6 +379,7 @@ func TestCategorize_ResolvesTenantModelInstanceCredentials(t *testing.T) { pushComponentDB(t, db) if err := db.Create(&entity.TenantModelProvider{ ID: "provider-1", + UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "SILICONFLOW", }).Error; err != nil { @@ -428,6 +429,7 @@ func TestCategorize_ResolvesInstanceCredentialsModelNameWithAt(t *testing.T) { pushComponentDB(t, db) if err := db.Create(&entity.TenantModelProvider{ ID: "provider-1", + UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "LM-Studio", }).Error; err != nil { @@ -478,6 +480,7 @@ func TestCategorize_ResolvesTenantModelID(t *testing.T) { pushComponentDB(t, db) if err := db.Create(&entity.TenantModelProvider{ ID: "provider-1", + UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "SILICONFLOW", }).Error; err != nil { @@ -538,6 +541,7 @@ func TestCategorize_ResolvesSoleActiveInstanceWhenDefaultMissing(t *testing.T) { pushComponentDB(t, db) if err := db.Create(&entity.TenantModelProvider{ ID: "provider-1", + UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "SILICONFLOW", }).Error; err != nil { diff --git a/internal/agent/component/llm.go b/internal/agent/component/llm.go index d64a5e21f8..adbe2a77d7 100644 --- a/internal/agent/component/llm.go +++ b/internal/agent/component/llm.go @@ -30,8 +30,8 @@ import ( "ragflow/internal/agent/component/prompts" "ragflow/internal/agent/runtime" "ragflow/internal/common" - "ragflow/internal/dao" "ragflow/internal/entity/models" + "ragflow/internal/modelcontext" "go.uber.org/zap" ) @@ -815,7 +815,7 @@ func (c *LLMComponent) prepareChatRequest( tenantID = tid } } - contentLength := dao.ResolveModelContentLength(ctx, db, tenantID, originalModelID, p.Driver, p.ModelID) + contentLength := modelcontext.ResolveModelContentLength(ctx, db, tenantID, originalModelID, p.Driver, p.ModelID) if contentLength <= 0 { // A 0 makes chat.FitMessages fall back to the 8192 default budget, which // can silently discard most of a large-context prompt, so surface the diff --git a/internal/agent/component/llm_credentials.go b/internal/agent/component/llm_credentials.go index f48b3ee547..9c7dcd3173 100644 --- a/internal/agent/component/llm_credentials.go +++ b/internal/agent/component/llm_credentials.go @@ -123,7 +123,13 @@ func resolveTenantModelInstanceCredentials(ctx context.Context, db *gorm.DB, tid zap.String("model", modelName), zap.String("instance", instanceName)) - provider, err := dao.NewTenantModelProviderDAO().GetByTenantIDAndProviderName(ctx, db, tid, providerName) + owner, ownerErr := dao.NewUserTenantDAO().GetOwnerByTenantID(ctx, db, tid) + if ownerErr != nil { + common.Debug("llm credentials: new-table fallback: tenant owner lookup failed", zap.String("tenant", tid), zap.Error(ownerErr)) + return "", baseURL, false + } + ownerUserID := owner.UserID + provider, err := dao.NewTenantModelProviderDAO().GetByTenantIDAndUserIDAndProviderName(ctx, db, tid, ownerUserID, providerName) if err != nil || provider == nil { common.Debug("llm credentials: new-table fallback: provider not found", zap.String("provider", providerName), zap.Error(err)) return "", baseURL, false diff --git a/internal/agent/component/llm_test.go b/internal/agent/component/llm_test.go index a334df4f49..8895b973be 100644 --- a/internal/agent/component/llm_test.go +++ b/internal/agent/component/llm_test.go @@ -275,6 +275,7 @@ func TestLLM_Invoke_CompositeModel_CustomContextOverride(t *testing.T) { if err := db.Create(&entity.TenantModelProvider{ ID: "provider-comp-1", + UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "OpenAI", }).Error; err != nil { @@ -337,6 +338,7 @@ func TestLLM_Invoke_UUIDModel_CustomContextOverride(t *testing.T) { if err := db.Create(&entity.TenantModelProvider{ ID: "provider-uuid-2", + UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "OpenAI", }).Error; err != nil { @@ -400,6 +402,7 @@ func TestLLM_Invoke_UUIDModel_ResolvesContentLength(t *testing.T) { if err := db.Create(&entity.TenantModelProvider{ ID: "provider-uuid-1", + UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "OpenAI", }).Error; err != nil { @@ -459,6 +462,7 @@ func TestLLM_ResolvesTenantModelID(t *testing.T) { if err := db.Create(&entity.TenantModelProvider{ ID: "provider-1", + UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "DeepSeek", }).Error; err != nil { @@ -577,7 +581,7 @@ func TestLLM_Invoke_UnresolvableModelFallsBackTo8192(t *testing.T) { // TestLLM_Invoke_UsesModelContentLengthBudget verifies that the message // fitting budget in Invoke is the chat model's context window -// (content_length) resolved via dao.ResolveModelContentLength — NOT the +// (content_length) resolved via modelcontext.ResolveModelContentLength - NOT the // canvas max_tokens / the 8192 fallback. A user prompt far larger than the // 8192 fallback (but well inside gpt-4o@openai's 128k window) must be passed // through to the invoker untrimmed. diff --git a/internal/dao/ingestion_parser_config_runtime_test.go b/internal/dao/ingestion_parser_config_runtime_test.go index b4b1625f15..e8a35704de 100644 --- a/internal/dao/ingestion_parser_config_runtime_test.go +++ b/internal/dao/ingestion_parser_config_runtime_test.go @@ -200,7 +200,7 @@ func TestMigratedMinerUControlsReachRequest(t *testing.T) { })) t.Cleanup(server.Close) db := openIngestionRuntimeDB(t) - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatal(err) } extra, err := json.Marshal(map[string]string{"base_url": server.URL}) @@ -208,7 +208,7 @@ func TestMigratedMinerUControlsReachRequest(t *testing.T) { t.Fatal(err) } for _, row := range []any{ - &entity.TenantModelProvider{ID: "provider-mineru", TenantID: "tenant", ProviderName: "MinerU"}, + &entity.TenantModelProvider{ID: "provider-mineru", UserID: "tenant", TenantID: "tenant", ProviderName: "MinerU"}, &entity.TenantModelInstance{ID: "instance-mineru", ProviderID: "provider-mineru", InstanceName: "default", APIKey: "synthetic-mineru-key", Status: "active", Extra: string(extra)}, &entity.TenantModel{ID: "model-mineru", ProviderID: "provider-mineru", InstanceID: "instance-mineru", ModelName: "MinerU-model", ModelType: int(entity.ModelTypeOCR), Status: "active"}, } { @@ -715,7 +715,7 @@ func openIngestionRuntimeDB(t *testing.T) *gorm.DB { t.Fatal(err) } t.Cleanup(func() { _ = sqlDB.Close() }) - if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.SystemSettings{}); err != nil { + if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.SystemSettings{}, &entity.UserTenant{}); err != nil { t.Fatal(err) } return db diff --git a/internal/dao/tenant_model.go b/internal/dao/tenant_model.go index cacdea2b05..2c0c4668f0 100644 --- a/internal/dao/tenant_model.go +++ b/internal/dao/tenant_model.go @@ -18,6 +18,7 @@ package dao import ( "context" + "fmt" "ragflow/internal/entity" "gorm.io/gorm" @@ -69,6 +70,24 @@ func (dao *TenantModelDAO) GetByIDs(ctx context.Context, db *gorm.DB, ids []stri return models, nil } +// ListByTenantIDAndModelType returns a tenant's authorized model records. +// Zero-type rows are retained when filtering so the caller can derive their +// type from the provider catalog before producing the response. +func (dao *TenantModelDAO) ListByTenantIDAndModelType(ctx context.Context, db *gorm.DB, tenantID string, modelType int, authorize func(*gorm.DB) *gorm.DB) ([]*entity.TenantModel, error) { + if authorize == nil { + return nil, fmt.Errorf("model list requires an authorization scope") + } + var models []*entity.TenantModel + query := db.WithContext(ctx).Model(&entity.TenantModel{}). + Joins("JOIN tenant_model_provider ON tenant_model_provider.id = tenant_model.provider_id"). + Where("tenant_model_provider.tenant_id = ?", tenantID) + if modelType != 0 { + query = query.Where("(tenant_model.model_type = 0 OR (tenant_model.model_type & ?) > 0)", modelType) + } + err := authorize(query).Find(&models).Error + return models, err +} + func (dao *TenantModelDAO) GetModelByProviderIDAndInstanceIDAndModelName(ctx context.Context, db *gorm.DB, providerID, instanceID, modelName string) (*entity.TenantModel, error) { var model entity.TenantModel err := db.WithContext(ctx).Where("provider_id = ? AND instance_id = ? AND model_name = ?", providerID, instanceID, modelName).First(&model).Error @@ -140,7 +159,7 @@ func (dao *TenantModelDAO) DeleteByInstanceIDs(ctx context.Context, db *gorm.DB, // used to fetch per-tenant enable/disable overrides in bulk during // /api/v1/models response assembly. The Go port never WRITES to // tenant_model, so callers must treat an empty result as "use factory -// defaults" — see ModelProviderService.ListTenantAddedModels. +// defaults" - see ModelProviderService.ListModelsForUserInTenant. func (dao *TenantModelDAO) GetModelsByProviderIDsAndInstanceIDs(ctx context.Context, db *gorm.DB, providerIDs, instanceIDs []string) ([]*entity.TenantModel, error) { models := make([]*entity.TenantModel, 0) if len(providerIDs) == 0 || len(instanceIDs) == 0 { diff --git a/internal/dao/tenant_model_instance.go b/internal/dao/tenant_model_instance.go index c71c2c86a8..a8afa475a3 100644 --- a/internal/dao/tenant_model_instance.go +++ b/internal/dao/tenant_model_instance.go @@ -28,6 +28,14 @@ import ( // TenantModelInstanceDAO tenant model instance data access object type TenantModelInstanceDAO struct{} +// TenantModelInstanceSummary contains the fields needed by model list views. +type TenantModelInstanceSummary struct { + ID string + InstanceName string + ProviderID string + Status string +} + // NewTenantModelInstanceDAO create tenant model instance DAO func NewTenantModelInstanceDAO() *TenantModelInstanceDAO { return &TenantModelInstanceDAO{} @@ -67,21 +75,30 @@ func (dao *TenantModelInstanceDAO) GetAllInstancesByProviderID(ctx context.Conte return instances, nil } -// GetByProviderIDs returns all TenantModelInstance rows whose provider_id -// is in providerIDs. Mirrors Python's -// TenantModelInstanceService.get_by_provider_ids used by -// models_api_service.list_tenant_added_models. An empty input slice -// returns an empty (non-nil) slice with no error. -func (dao *TenantModelInstanceDAO) GetByProviderIDs(ctx context.Context, db *gorm.DB, providerIDs []string) ([]*entity.TenantModelInstance, error) { - instances := make([]*entity.TenantModelInstance, 0) +// ListSummariesByProviderIDs loads only non-secret instance fields used by model lists. +func (dao *TenantModelInstanceDAO) ListSummariesByProviderIDs(ctx context.Context, db *gorm.DB, providerIDs []string) ([]*TenantModelInstanceSummary, error) { + instances := make([]*TenantModelInstanceSummary, 0) if len(providerIDs) == 0 { return instances, nil } - err := db.WithContext(ctx).Where("provider_id IN ?", providerIDs).Find(&instances).Error - if err != nil { - return nil, err + err := db.WithContext(ctx).Model(&entity.TenantModelInstance{}). + Select("id", "instance_name", "provider_id", "status"). + Where("provider_id IN ?", providerIDs). + Find(&instances).Error + return instances, err +} + +// ListProviderIDsWithInstances returns provider IDs that have at least one instance. +func (dao *TenantModelInstanceDAO) ListProviderIDsWithInstances(ctx context.Context, db *gorm.DB, providerIDs []string) ([]string, error) { + instances := make([]string, 0) + if len(providerIDs) == 0 { + return instances, nil } - return instances, nil + err := db.WithContext(ctx).Model(&entity.TenantModelInstance{}). + Distinct("provider_id"). + Where("provider_id IN ?", providerIDs). + Pluck("provider_id", &instances).Error + return instances, err } func (dao *TenantModelInstanceDAO) GetInstanceByApiKey(ctx context.Context, db *gorm.DB, apiKey, providerID string) (*entity.TenantModelInstance, error) { diff --git a/internal/dao/tenant_model_max_tokens_migration_test.go b/internal/dao/tenant_model_max_tokens_migration_test.go index 16dc63639f..eff4ee194b 100644 --- a/internal/dao/tenant_model_max_tokens_migration_test.go +++ b/internal/dao/tenant_model_max_tokens_migration_test.go @@ -26,6 +26,16 @@ import ( "ragflow/internal/entity" ) +func migrationTestMaxTokens(extra string) int64 { + var value struct { + MaxTokens int64 `json:"max_tokens"` + } + if err := json.Unmarshal([]byte(extra), &value); err != nil { + return 0 + } + return value.MaxTokens +} + func setupTenantModelMaxTokensMigrationDB(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true}) @@ -154,7 +164,7 @@ func TestMigrateTenantModelMaxTokensBackfillsOnlyMissingValues(t *testing.T) { t.Fatalf("failed to read tenant model %q: %v", item.id, err) } } - if got := modelExtraMaxTokens(missing.Extra); got != 65536 { + if got := migrationTestMaxTokens(missing.Extra); got != 65536 { t.Fatalf("backfilled max_tokens = %d, want 65536 (extra: %s)", got, missing.Extra) } var missingExtra struct { @@ -166,10 +176,10 @@ func TestMigrateTenantModelMaxTokensBackfillsOnlyMissingValues(t *testing.T) { if !missingExtra.IsTools { t.Fatalf("backfill lost existing extra fields: %s", missing.Extra) } - if got := modelExtraMaxTokens(configured.Extra); got != 16384 { + if got := migrationTestMaxTokens(configured.Extra); got != 16384 { t.Fatalf("existing max_tokens was overwritten: got %d, want 16384", got) } - if got := modelExtraMaxTokens(zero.Extra); got != 0 { + if got := migrationTestMaxTokens(zero.Extra); got != 0 { t.Fatalf("explicit zero max_tokens was overwritten: got %d, want 0", got) } if hasMaxTokens, err := tenantModelExtraHasMaxTokens(zero.Extra); err != nil || !hasMaxTokens { @@ -239,7 +249,7 @@ func TestBackfillTenantModelMaxTokensSkipsCatalogDefaults(t *testing.T) { if tc.wantUpdated == 0 && model.Extra != `{"is_tools":true}` { t.Fatalf("catalog default changed extra: %s", model.Extra) } - if tc.wantUpdated == 1 && int64(modelExtraMaxTokens(model.Extra)) != tc.limit { + if tc.wantUpdated == 1 && migrationTestMaxTokens(model.Extra) != tc.limit { t.Fatalf("custom limit was not preserved: %s", model.Extra) } }) diff --git a/internal/dao/tenant_model_provider.go b/internal/dao/tenant_model_provider.go index cc06228d59..33a167e7bc 100644 --- a/internal/dao/tenant_model_provider.go +++ b/internal/dao/tenant_model_provider.go @@ -45,33 +45,29 @@ func (dao *TenantModelProviderDAO) GetByID(ctx context.Context, db *gorm.DB, id return &provider, nil } -// GetByTenantIDAndProviderName get the providers by tenant ID and provider name -func (dao *TenantModelProviderDAO) GetByTenantIDAndProviderName(ctx context.Context, db *gorm.DB, tenantID, providerName string) (*entity.TenantModelProvider, error) { +// GetByTenantIDAndUserIDAndProviderName gets a user's provider configuration in a tenant. +func (dao *TenantModelProviderDAO) GetByTenantIDAndUserIDAndProviderName(ctx context.Context, db *gorm.DB, tenantID, userID, providerName string) (*entity.TenantModelProvider, error) { var provider entity.TenantModelProvider - err := db.WithContext(ctx).Where("tenant_id = ? AND provider_name = ?", tenantID, providerName).First(&provider).Error + err := db.WithContext(ctx). + Where("tenant_id = ? AND user_id = ? AND LOWER(provider_name) = LOWER(?)", tenantID, userID, providerName). + First(&provider).Error if err != nil { return nil, err } return &provider, nil } -// DeleteByTenantIDAndProviderName DeleteByTenantID deletes all providers by tenant ID (hard delete) -func (dao *TenantModelProviderDAO) DeleteByTenantIDAndProviderName(ctx context.Context, db *gorm.DB, tenantID, providerName string) (int64, error) { - result := db.WithContext(ctx).Unscoped().Where("tenant_id = ? AND provider_name = ?", tenantID, providerName).Delete(&entity.TenantModelProvider{}) - return result.RowsAffected, result.Error +// ListByTenantIDAndUserID returns one user's provider configurations in a tenant. +func (dao *TenantModelProviderDAO) ListByTenantIDAndUserID(ctx context.Context, db *gorm.DB, tenantID, userID string) ([]*entity.TenantModelProvider, error) { + var providers []*entity.TenantModelProvider + err := db.WithContext(ctx). + Where("tenant_id = ? AND user_id = ?", tenantID, userID). + Find(&providers).Error + return providers, err } -// GetByTenantID returns all TenantModelProvider rows for a tenant. -// Mirrors Python's TenantModelProviderService.get_by_tenant_id and is the -// entry point for /api/v1/models ("list all added models"). The Go port -// uses this to enumerate which providers a tenant has linked before -// fanning out to TenantModelInstanceDAO / TenantModelDAO for the joined -// result. -func (dao *TenantModelProviderDAO) GetByTenantID(ctx context.Context, db *gorm.DB, tenantID string) ([]*entity.TenantModelProvider, error) { - var providers []*entity.TenantModelProvider - err := db.WithContext(ctx).Where("tenant_id = ?", tenantID).Find(&providers).Error - if err != nil { - return nil, err - } - return providers, nil +// DeleteByID hard-deletes one provider configuration. +func (dao *TenantModelProviderDAO) DeleteByID(ctx context.Context, db *gorm.DB, id string) (int64, error) { + result := db.WithContext(ctx).Unscoped().Where("id = ?", id).Delete(&entity.TenantModelProvider{}) + return result.RowsAffected, result.Error } diff --git a/internal/dao/user_tenant.go b/internal/dao/user_tenant.go index bd5360a25b..c53c7dad83 100644 --- a/internal/dao/user_tenant.go +++ b/internal/dao/user_tenant.go @@ -129,6 +129,9 @@ func (dao *UserTenantDAO) GetOwnerByTenantID(ctx context.Context, db *gorm.DB, t Where("tenant_id = ? AND user_id = ? AND status = ?", tenantID, tenantID, entity.TenantMembershipStatusJoined). Take(&membership).Error if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return &entity.UserTenant{UserID: tenantID, TenantID: tenantID}, nil + } return nil, err } return &membership, nil diff --git a/internal/dao/user_tenant_test.go b/internal/dao/user_tenant_test.go index 42485db432..2f12ed912e 100644 --- a/internal/dao/user_tenant_test.go +++ b/internal/dao/user_tenant_test.go @@ -50,3 +50,21 @@ func TestGetOwnerByTenantIDFallsBackToTenantIDUserID(t *testing.T) { t.Fatalf("GetOwnerByTenantID() = %#v, want the tenant_id=user_id membership", membership) } } + +func TestGetOwnerByTenantIDUsesTenantIDWhenOwnerRowIsMissing(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&entity.UserTenant{}); err != nil { + t.Fatalf("migrate user_tenant: %v", err) + } + + owner, err := NewUserTenantDAO().GetOwnerByTenantID(t.Context(), db, "personal-tenant") + if err != nil { + t.Fatalf("GetOwnerByTenantID() error = %v", err) + } + if owner.UserID != "personal-tenant" || owner.TenantID != "personal-tenant" { + t.Fatalf("GetOwnerByTenantID() = %#v, want tenant ID as the personal owner", owner) + } +} diff --git a/internal/handler/chat.go b/internal/handler/chat.go index 532db75869..c8039f72a4 100644 --- a/internal/handler/chat.go +++ b/internal/handler/chat.go @@ -189,7 +189,7 @@ func (h *ChatHandler) Create(c *gin.Context) { // @Success 200 {object} map[string]interface{} // @Router /api/v1/chat/mindmap [post] func (h *ChatHandler) MindMap(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return @@ -207,13 +207,13 @@ func (h *ChatHandler) MindMap(c *gin.Context) { ctx := c.Request.Context() searchConfig := map[string]interface{}{} - modelTenantID := user.ID + var modelTenantID string if req.SearchID != "" { if h.searchSvc == nil { jsonInternalError(c, fmt.Errorf("search service not configured")) return } - detail, err := h.searchSvc.GetDetail(ctx, user.ID, req.SearchID) + detail, err := h.searchSvc.GetDetail(ctx, req.SearchID) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -238,7 +238,6 @@ func (h *ChatHandler) MindMap(c *gin.Context) { KbIDs: kbIDs, SearchID: req.SearchID, SearchConfig: searchConfig, - AuthUserID: user.ID, ModelTenantID: modelTenantID, ChunkSvc: h.chunkSvc, LLM: h.modelFactory, diff --git a/internal/handler/chat_audio.go b/internal/handler/chat_audio.go index a00aa5cdaf..6595d013f6 100644 --- a/internal/handler/chat_audio.go +++ b/internal/handler/chat_audio.go @@ -73,7 +73,7 @@ func (h *ChatHandler) ChatAudioSpeech(c *gin.Context) { return } - ttsModel, err := h.modelFactory.NewDefaultTTSModel(ctx, service.ModelAccess{UserID: user.ID, TenantID: user.ID}) + ttsModel, err := h.modelFactory.NewDefaultTTSModel(ctx, service.Identity{UserID: user.ID, TenantID: user.ID}) if err != nil { common.ErrorWithCode(c, common.CodeDataError, err.Error()) return @@ -268,7 +268,7 @@ func (h *ChatHandler) ChatAudioTranscription(c *gin.Context) { return } - asrModel, err := h.modelFactory.NewDefaultASRModel(ctx, service.ModelAccess{UserID: user.ID, TenantID: user.ID}) + asrModel, err := h.modelFactory.NewDefaultASRModel(ctx, service.Identity{UserID: user.ID, TenantID: user.ID}) if err != nil { common.ErrorWithCode(c, common.CodeDataError, err.Error()) return diff --git a/internal/handler/chat_recommendation.go b/internal/handler/chat_recommendation.go index 19c79cd0f9..aac3fbecce 100644 --- a/internal/handler/chat_recommendation.go +++ b/internal/handler/chat_recommendation.go @@ -43,7 +43,7 @@ type ChatRecommendationRequest struct { // @Success 200 {object} map[string]interface{} // @Router /api/v1/chat/recommendation [post] func (h *ChatHandler) Recommendation(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return @@ -59,7 +59,7 @@ func (h *ChatHandler) Recommendation(c *gin.Context) { return } ctx := c.Request.Context() - questions, err := service.GenerateRelatedQuestions(ctx, user.ID, req.Question, req.SearchID, h.searchSvc, h.tenantSvc, h.modelFactory) + questions, err := service.GenerateRelatedQuestions(ctx, req.Question, req.SearchID, h.searchSvc, h.modelFactory) if err != nil { common.Warn("chat recommendation failed", zap.String("error", err.Error())) common.ResponseWithCodeData(c, common.CodeOperatingError, nil, err.Error()) diff --git a/internal/handler/compilation_template.go b/internal/handler/compilation_template.go index ee4f8adb46..e37e2270cd 100644 --- a/internal/handler/compilation_template.go +++ b/internal/handler/compilation_template.go @@ -43,12 +43,12 @@ func NewCompilationTemplateHandler(svc *service.CompilationTemplateService) *Com // @Success 200 {object} map[string]interface{} // @Router /v1/compilation-templates/builtins [get] func (h *CompilationTemplateHandler) ListBuiltins(c *gin.Context) { - user, code, msg := GetUser(c) + _, code, msg := GetUser(c) if code != common.CodeSuccess { common.ErrorWithCode(c, code, msg) return } - templates, err := h.compilationTemplateService.ListBuiltins(c.Request.Context(), user.ID) + templates, err := h.compilationTemplateService.ListBuiltins(c.Request.Context()) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, nil, err.Error()) return diff --git a/internal/handler/dify_retrieval_handler.go b/internal/handler/dify_retrieval_handler.go index 4e1c8383e6..5bee1ecb0c 100644 --- a/internal/handler/dify_retrieval_handler.go +++ b/internal/handler/dify_retrieval_handler.go @@ -50,8 +50,8 @@ type KBServiceIface interface { // ModelCreator supplies typed model instances to Dify retrieval. type ModelCreator interface { - NewEmbeddingModel(ctx context.Context, access service.ModelAccess, modelRef string) (*modelModule.EmbeddingModel, error) - NewDefaultChatModel(ctx context.Context, access service.ModelAccess) (*modelModule.ChatModel, error) + NewEmbeddingModel(ctx context.Context, access service.Identity, modelRef string) (*modelModule.EmbeddingModel, error) + NewDefaultChatModel(ctx context.Context, access service.Identity) (*modelModule.ChatModel, error) } // MetadataServiceIface abstracts MetadataService for the Dify handler. @@ -240,7 +240,7 @@ func (h *DifyRetrievalHandler) Retrieval(c *gin.Context) { } // Get embedding model - access := service.ModelAccess{UserID: user.ID, TenantID: kb.TenantID} + access := service.Identity{UserID: user.ID, TenantID: kb.TenantID} embModel, err := h.modelFactory.NewEmbeddingModel(ctx, access, kb.EmbdID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, fmt.Sprintf("failed to get embedding model: %v", err)) diff --git a/internal/handler/dify_retrieval_handler_test.go b/internal/handler/dify_retrieval_handler_test.go index c00b5cc936..38517b95eb 100644 --- a/internal/handler/dify_retrieval_handler_test.go +++ b/internal/handler/dify_retrieval_handler_test.go @@ -65,12 +65,12 @@ func (m *mockKBService) CheckAccess(ctx context.Context, subject permission.Subj type mockModelFactory struct{} -func (mockModelFactory) NewEmbeddingModel(context.Context, service.ModelAccess, string) (*modelModule.EmbeddingModel, error) { +func (mockModelFactory) NewEmbeddingModel(context.Context, service.Identity, string) (*modelModule.EmbeddingModel, error) { name := "test-model" return modelModule.NewEmbeddingModel(&modelModule.DummyModel{}, &name, &modelModule.APIConfig{}, 0), nil } -func (mockModelFactory) NewDefaultChatModel(context.Context, service.ModelAccess) (*modelModule.ChatModel, error) { +func (mockModelFactory) NewDefaultChatModel(context.Context, service.Identity) (*modelModule.ChatModel, error) { name := "test-model" return modelModule.NewChatModel(&modelModule.DummyModel{}, &name, &modelModule.APIConfig{}), nil } diff --git a/internal/handler/mindmap.go b/internal/handler/mindmap.go index 679bbb6b7e..9fdad99743 100644 --- a/internal/handler/mindmap.go +++ b/internal/handler/mindmap.go @@ -34,7 +34,6 @@ type mindMapRunConfig struct { KbIDs common.StringSlice SearchID string SearchConfig map[string]interface{} - AuthUserID string ModelTenantID string ChunkSvc service.Retriever LLM *service.ModelFactory @@ -48,12 +47,16 @@ func runMindMap(ctx context.Context, config mindMapRunConfig) (mindMapNode, erro if config.LLM == nil { return mindMapNode{}, fmt.Errorf("LLM not configured") } + identity, err := service.ResolveIdentity(ctx) + if err != nil { + return mindMapNode{}, err + } modelTenantID := config.ModelTenantID if modelTenantID == "" { - modelTenantID = config.AuthUserID + modelTenantID = identity.TenantID } retrievalReq := mindMapRetrievalRequest(config.Question, config.KbIDs, config.SearchID, config.SearchConfig) - ranks, err := config.ChunkSvc.RetrievalTest(ctx, retrievalReq, config.AuthUserID) + ranks, err := config.ChunkSvc.RetrievalTest(ctx, retrievalReq) if err != nil { return mindMapNode{}, err } @@ -69,7 +72,8 @@ func runMindMap(ctx context.Context, config mindMapRunConfig) (mindMapNode, erro // A stale search_config chat_id should not prevent mind-map generation; // fall back to the tenant's configured default model. - access := service.ModelAccess{UserID: config.AuthUserID, TenantID: modelTenantID} + identity.TenantID = modelTenantID + access := identity chatModel, modelErr := config.LLM.NewChatModel(streamCtx, access, modelID) if modelErr != nil { chatModel, modelErr = config.LLM.NewDefaultChatModel(streamCtx, access) diff --git a/internal/handler/providers.go b/internal/handler/providers.go index 6b0049ef64..f5147a0135 100644 --- a/internal/handler/providers.go +++ b/internal/handler/providers.go @@ -1058,10 +1058,8 @@ func (h *ProviderHandler) ChatToModel(c *gin.Context) { Verbosity: req.Verbosity, } - userID := c.GetString("user_id") email := c.GetString("email") modelUsage := common.ModelUsage{ - UserID: userID, UserEmail: email, Type: "chat", StartAt: time.Now(), @@ -1117,7 +1115,6 @@ func (h *ProviderHandler) ChatToModel(c *gin.Context) { errorCode, err := h.modelCallService.ChatToModelStreamWithSender( ctx, modelRef, - userID, messages, &chatConfig, &modelUsage, @@ -1144,7 +1141,6 @@ func (h *ProviderHandler) ChatToModel(c *gin.Context) { response, errorCode, err = h.modelCallService.ChatToModelWithMessages( ctx, modelRef, - userID, messages, &chatConfig, &modelUsage, @@ -1203,7 +1199,6 @@ func (h *ProviderHandler) EmbedText(c *gin.Context) { } } - userID := c.GetString("user_id") modelRef := "" if req.ModelID != nil { modelRef = *req.ModelID @@ -1220,7 +1215,7 @@ func (h *ProviderHandler) EmbedText(c *gin.Context) { var errorCode common.ErrorCode var err error - response, errorCode, err = h.modelCallService.EmbedText(ctx, modelRef, userID, req.Texts, &embeddingConfig) + response, errorCode, err = h.modelCallService.EmbedText(ctx, modelRef, req.Texts, &embeddingConfig) if err != nil { common.ErrorWithCode(c, errorCode, err.Error()) return @@ -1270,7 +1265,6 @@ func (h *ProviderHandler) RerankDocument(c *gin.Context) { } } - userID := c.GetString("user_id") modelRef := "" if req.ModelID != nil { modelRef = *req.ModelID @@ -1287,7 +1281,7 @@ func (h *ProviderHandler) RerankDocument(c *gin.Context) { var errorCode common.ErrorCode var err error - response, errorCode, err = h.modelCallService.RerankDocument(ctx, modelRef, userID, req.Query, req.Documents, &rerankConfig) + response, errorCode, err = h.modelCallService.RerankDocument(ctx, modelRef, req.Query, req.Documents, &rerankConfig) if err != nil { common.ErrorWithCode(c, errorCode, err.Error()) return @@ -1338,7 +1332,6 @@ func (h *ProviderHandler) TranscribeAudio(c *gin.Context) { } } - userID := c.GetString("user_id") modelRef := "" if req.ModelID != nil { modelRef = *req.ModelID @@ -1385,7 +1378,7 @@ func (h *ProviderHandler) TranscribeAudio(c *gin.Context) { } // Stream response using sender function ( the best performance, no channel) - errorCode, err := h.modelCallService.TranscribeAudioStream(ctx, modelRef, userID, req.File, &asrConfig, sender) + errorCode, err := h.modelCallService.TranscribeAudioStream(ctx, modelRef, req.File, &asrConfig, sender) if errorCode != common.CodeSuccess { c.SSEvent("error", err.Error()) } @@ -1397,7 +1390,7 @@ func (h *ProviderHandler) TranscribeAudio(c *gin.Context) { var errorCode common.ErrorCode var err error - response, errorCode, err = h.modelCallService.TranscribeAudio(ctx, modelRef, userID, req.File, &asrConfig) + response, errorCode, err = h.modelCallService.TranscribeAudio(ctx, modelRef, req.File, &asrConfig) if err != nil { common.ErrorWithCode(c, errorCode, err.Error()) return @@ -1447,7 +1440,6 @@ func (h *ProviderHandler) AudioSpeech(c *gin.Context) { } } - userID := c.GetString("user_id") modelRef := "" if req.ModelID != nil { modelRef = *req.ModelID @@ -1494,7 +1486,7 @@ func (h *ProviderHandler) AudioSpeech(c *gin.Context) { } // Stream response using sender function ( the best performance, no channel) - errorCode, err := h.modelCallService.AudioSpeechStream(ctx, modelRef, userID, req.Text, &ttsConfig, sender) + errorCode, err := h.modelCallService.AudioSpeechStream(ctx, modelRef, req.Text, &ttsConfig, sender) if errorCode != common.CodeSuccess { c.SSEvent("error", err.Error()) } @@ -1506,7 +1498,7 @@ func (h *ProviderHandler) AudioSpeech(c *gin.Context) { var errorCode common.ErrorCode var err error - response, errorCode, err = h.modelCallService.AudioSpeech(ctx, modelRef, userID, req.Text, &ttsConfig) + response, errorCode, err = h.modelCallService.AudioSpeech(ctx, modelRef, req.Text, &ttsConfig) if err != nil { common.ErrorWithCode(c, errorCode, err.Error()) return @@ -1555,8 +1547,6 @@ func (h *ProviderHandler) OCRFile(c *gin.Context) { } } - userID := c.GetString("user_id") - modelRef := "" if req.ModelID != nil { modelRef = *req.ModelID @@ -1571,7 +1561,7 @@ func (h *ProviderHandler) OCRFile(c *gin.Context) { var errorCode common.ErrorCode var err error - response, errorCode, err = h.modelCallService.OCRFile(ctx, modelRef, userID, req.Content, req.URL, &OCRConfig) + response, errorCode, err = h.modelCallService.OCRFile(ctx, modelRef, req.Content, req.URL, &OCRConfig) if err != nil { common.ErrorWithCode(c, errorCode, err.Error()) return @@ -1623,8 +1613,6 @@ func (h *ProviderHandler) ParseFile(c *gin.Context) { } } - userID := c.GetString("user_id") - modelRef := "" if req.ModelID != nil { modelRef = *req.ModelID @@ -1645,7 +1633,7 @@ func (h *ProviderHandler) ParseFile(c *gin.Context) { var errorCode common.ErrorCode var err error - response, errorCode, err = h.modelCallService.ParseFile(ctx, modelRef, userID, req.Content, req.URL, &parseFileConfig) + response, errorCode, err = h.modelCallService.ParseFile(ctx, modelRef, req.Content, req.URL, &parseFileConfig) if err != nil { common.ErrorWithCode(c, errorCode, err.Error()) return @@ -1654,22 +1642,9 @@ func (h *ProviderHandler) ParseFile(c *gin.Context) { common.SuccessWithData(c, response, "success") } -// ListTenantAddedModels is the response handler for GET /api/v1/models. -// It is the Go port of Python's -// api/apps/restful_apis/models_api.py:get_added_models and feeds -// web/src/hooks/use-llm-request.tsx → useFetchAllAddedModels. The data -// shape is the array form (one row per (provider × instance × llm) with -// model_type: string[]), matching the IAddedModel interface in -// web/src/interfaces/database/llm.ts:64-71. -// -// The previous contract routed this path to TenantHandler.GetModels → -// TenantService.ListTenantDefaultModels, which only enumerates the 6-7 -// default tenant fields and returned `[]` for any tenant without -// defaults, breaking the front-end's "View Models" list. The Go port -// has no writers for tenant_model, so this endpoint must be driven by -// the factory catalog cross-referenced with the tenant's instance list — -// see service.ModelProviderService.ListTenantAddedModels. -func (h *ProviderHandler) ListTenantAddedModels(c *gin.Context) { +// ListModelsForUserInTenant handles GET /api/v1/models for the authenticated +// user and selected tenant context. +func (h *ProviderHandler) ListModelsForUserInTenant(c *gin.Context) { user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) @@ -1680,11 +1655,11 @@ func (h *ProviderHandler) ListTenantAddedModels(c *gin.Context) { tenantID := c.Query("tenant_id") ctx := c.Request.Context() - addedModels, code, err := h.modelProviderService.ListTenantAddedModels(ctx, user.ID, tenantID, modelType) + models, code, err := h.modelProviderService.ListModelsForUserInTenant(ctx, user.ID, tenantID, modelType) if err != nil { common.ErrorWithCode(c, code, err.Error()) return } - common.SuccessWithData(c, addedModels, "success") + common.SuccessWithData(c, models, "success") } diff --git a/internal/handler/search.go b/internal/handler/search.go index 3020eba3a9..0b976c266b 100644 --- a/internal/handler/search.go +++ b/internal/handler/search.go @@ -85,12 +85,11 @@ func getOwnerIDs(c *gin.Context) []string { // @Success 200 {object} service.ListSearchAppsResponse // @Router /api/v1/searches [get] func (h *SearchHandler) ListSearches(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return } - userID := user.ID // Parse query parameters keywords := c.Query("keywords") @@ -132,7 +131,7 @@ func (h *SearchHandler) ListSearches(c *gin.Context) { // List search apps with filtering ctx := c.Request.Context() - result, err := h.searchService.ListSearches(ctx, userID, keywords, page, pageSize, terms, ownerIDs) + result, err := h.searchService.ListSearches(ctx, keywords, page, pageSize, terms, ownerIDs) if err != nil { if respondPermissionErrorIf(c, err, false) { return @@ -169,7 +168,12 @@ func (h *SearchHandler) CreateSearch(c *gin.Context) { return } - userID := user.ID + // Get current user from context (same as Python current_user) + _, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + common.ErrorWithCode(c, errorCode, errorMessage) + return + } // Parse request body (same as Python get_request_json()) var req CreateSearchRequest @@ -185,7 +189,7 @@ func (h *SearchHandler) CreateSearch(c *gin.Context) { // Create search (same as Python SearchService.save within DB.atomic()) ctx := c.Request.Context() - result, err := h.searchService.CreateSearch(ctx, userID, req.Name, req.Description) + result, err := h.searchService.CreateSearch(ctx, req.Name, req.Description) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeBadRequest, nil, err.Error()) return @@ -206,12 +210,11 @@ func (h *SearchHandler) CreateSearch(c *gin.Context) { // @Router /api/v1/searches/{search_id} [get] func (h *SearchHandler) GetSearch(c *gin.Context) { // Get current user from context (same as Python current_user) - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return } - userID := user.ID // Get search_id from path parameter (same as Python ) searchID := c.Param("search_id") @@ -222,7 +225,7 @@ func (h *SearchHandler) GetSearch(c *gin.Context) { // Get search detail with permission check ctx := c.Request.Context() - search, err := h.searchService.GetSearchDetail(ctx, userID, searchID) + search, err := h.searchService.GetSearchDetail(ctx, searchID) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -263,12 +266,11 @@ func (h *SearchHandler) GetSearch(c *gin.Context) { // @Router /api/v1/searches/{search_id} [delete] func (h *SearchHandler) DeleteSearch(c *gin.Context) { // Get current user from context (same as Python current_user) - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return } - userID := user.ID // Get search_id from path parameter (same as Python ) searchID := c.Param("search_id") @@ -279,7 +281,7 @@ func (h *SearchHandler) DeleteSearch(c *gin.Context) { // Delete search with permission check ctx := c.Request.Context() - err := h.searchService.DeleteSearch(ctx, userID, searchID) + err := h.searchService.DeleteSearch(ctx, searchID) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -305,12 +307,11 @@ func (h *SearchHandler) DeleteSearch(c *gin.Context) { // @Router /api/v1/searches/{search_id} [put] func (h *SearchHandler) UpdateSearch(c *gin.Context) { // Get current user from context (same as Python current_user) - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return } - userID := user.ID // Get search_id from path parameter (same as Python ) searchID := c.Param("search_id") @@ -334,7 +335,7 @@ func (h *SearchHandler) UpdateSearch(c *gin.Context) { // Update search ctx := c.Request.Context() - updatedSearch, err := h.searchService.UpdateSearch(ctx, userID, searchID, &req) + updatedSearch, err := h.searchService.UpdateSearch(ctx, searchID, &req) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -376,7 +377,7 @@ func (h *SearchHandler) UpdateSearch(c *gin.Context) { } func (h *SearchHandler) Completion(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return @@ -394,7 +395,7 @@ func (h *SearchHandler) Completion(c *gin.Context) { searchSvc = service.NewSearchService() } ctx := c.Request.Context() - plan, code, err := searchSvc.PrepareCompletion(ctx, user.ID, c.Param("search_id"), &req) + plan, code, err := searchSvc.PrepareCompletion(ctx, c.Param("search_id"), &req) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -434,14 +435,14 @@ func (h *SearchHandler) Completion(c *gin.Context) { return } - adapter := &service.TenantStreamAdapter{Factory: h.modelFactory, TenantID: plan.UserID, ModelID: plan.ModelID} + adapter := &service.TenantStreamAdapter{Factory: h.modelFactory, Identity: plan.Identity, ModelID: plan.ModelID} ctx, cancel := context.WithCancel(c.Request.Context()) defer cancel() c.Request = c.Request.WithContext(ctx) hadError := false - for delta := range h.askService.StreamWithOptions(ctx, adapter, plan.UserID, plan.Question, plan.DatasetIDs, plan.Options) { + for delta := range h.askService.StreamWithOptions(ctx, adapter, plan.Question, plan.DatasetIDs, plan.Options) { var ok bool switch delta.Kind { case service.AskDeltaAnswer: diff --git a/internal/handler/searchbot.go b/internal/handler/searchbot.go index ae20d85891..266f9a878e 100644 --- a/internal/handler/searchbot.go +++ b/internal/handler/searchbot.go @@ -131,7 +131,7 @@ func (h *SearchBotHandler) SetAskService(svc *service.AskService) { h.askSvc = s // @Router /api/v1/searchbots/related-questions [post] func (h *SearchBotHandler) Handle(c *gin.Context) { ctx := c.Request.Context() - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return @@ -148,7 +148,7 @@ func (h *SearchBotHandler) Handle(c *gin.Context) { return } - questions, err := service.GenerateRelatedQuestions(ctx, user.ID, req.Question, req.SearchID, h.searchSvc, h.tenantSvc, h.modelFactory) + questions, err := service.GenerateRelatedQuestions(ctx, req.Question, req.SearchID, h.searchSvc, h.modelFactory) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -169,7 +169,7 @@ func (h *SearchBotHandler) Handle(c *gin.Context) { // @Success 200 {object} map[string]interface{} // @Router /api/v1/searchbots/retrieval [post] func (h *SearchBotHandler) RetrievalTest(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ResponseWithHttpCodeData(c, http.StatusUnauthorized, errorCode, nil, errorMessage) return @@ -211,7 +211,7 @@ func (h *SearchBotHandler) RetrievalTest(c *gin.Context) { svcReq := toRetrievalServiceRequest(&req) ctx := c.Request.Context() - result, err := h.chunkSvc.RetrievalTest(ctx, svcReq, user.ID) + result, err := h.chunkSvc.RetrievalTest(ctx, svcReq) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -234,7 +234,7 @@ func (h *SearchBotHandler) RetrievalTest(c *gin.Context) { // @Success 200 {object} map[string]interface{} // @Router /api/v1/searchbots/ask [post] func (h *SearchBotHandler) Ask(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return @@ -269,13 +269,21 @@ func (h *SearchBotHandler) Ask(c *gin.Context) { // Resolve chat model ID. modelID := "" options := service.AskStreamOptions{} + askAccess, accessErr := service.ResolveIdentity(c.Request.Context()) + if accessErr != nil { + if respondPermissionErrorIf(c, accessErr, true) { + return + } + jsonInternalError(c, accessErr) + return + } if req.SearchID != "" { if h.searchSvc == nil { jsonInternalError(c, fmt.Errorf("search service not configured")) return } ctx := c.Request.Context() - detail, err := h.searchSvc.GetDetail(ctx, user.ID, req.SearchID) + detail, err := h.searchSvc.GetDetail(ctx, req.SearchID) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -289,11 +297,10 @@ func (h *SearchBotHandler) Ask(c *gin.Context) { modelID = chatID } } - if modelID == "" && h.tenantSvc != nil { + if modelID == "" && h.modelFactory != nil { ctx := c.Request.Context() - defaultModel, err := h.tenantSvc.GetDefaultModelNameForUser(ctx, user.ID, entity.ModelTypeChat) - if err == nil && defaultModel != "" { - modelID = defaultModel + if info, infoErr := h.modelFactory.ResolveDefaultInfo(ctx, askAccess, entity.ModelTypeChat); infoErr == nil { + modelID = service.ModelReference(info) } } if modelID == "" { @@ -318,8 +325,8 @@ func (h *SearchBotHandler) Ask(c *gin.Context) { defer cancel() c.Request = c.Request.WithContext(ctx) hadError := false - adapter := &service.TenantStreamAdapter{Factory: h.modelFactory, TenantID: user.ID, ModelID: modelID} - for delta := range h.askSvc.StreamWithOptions(ctx, adapter, user.ID, req.Question, filtered, options) { + adapter := &service.TenantStreamAdapter{Factory: h.modelFactory, Identity: askAccess, ModelID: modelID} + for delta := range h.askSvc.StreamWithOptions(ctx, adapter, req.Question, filtered, options) { var ok bool switch delta.Kind { case service.AskDeltaAnswer: @@ -353,7 +360,7 @@ func (h *SearchBotHandler) Ask(c *gin.Context) { // @Success 200 {object} map[string]interface{} // @Router /api/v1/searchbots/mindmap [post] func (h *SearchBotHandler) MindMap(c *gin.Context) { - user, errorCode, errorMessage := GetUser(c) + _, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { common.ErrorWithCode(c, errorCode, errorMessage) return @@ -391,7 +398,7 @@ func (h *SearchBotHandler) MindMap(c *gin.Context) { return } ctx := c.Request.Context() - detail, err := h.searchSvc.GetDetail(ctx, user.ID, req.SearchID) + detail, err := h.searchSvc.GetDetail(ctx, req.SearchID) if err != nil { if respondPermissionErrorIf(c, err, true) { return @@ -404,15 +411,13 @@ func (h *SearchBotHandler) MindMap(c *gin.Context) { ctx := c.Request.Context() mindMap, err := runMindMap(ctx, mindMapRunConfig{ - Question: req.Question, - KbIDs: filtered, - SearchID: req.SearchID, - SearchConfig: searchConfig, - AuthUserID: user.ID, - ModelTenantID: user.ID, - ChunkSvc: h.chunkSvc, - LLM: h.modelFactory, - TenantSvc: h.tenantSvc, + Question: req.Question, + KbIDs: filtered, + SearchID: req.SearchID, + SearchConfig: searchConfig, + ChunkSvc: h.chunkSvc, + LLM: h.modelFactory, + TenantSvc: h.tenantSvc, }) if err != nil { if respondPermissionErrorIf(c, err, true) { @@ -449,7 +454,7 @@ func (h *SearchBotHandler) SearchBotDetail(c *gin.Context) { jsonInternalError(c, fmt.Errorf("search service not configured")) return } - detail, err := h.searchSvc.GetSearchShareDetail(ctx, identity.User.ID, searchID) + detail, err := h.searchSvc.GetSearchShareDetail(ctx, searchID) if err != nil { if respondPermissionErrorIf(c, err, true) { return diff --git a/internal/handler/searchbot_test.go b/internal/handler/searchbot_test.go index e815a21731..7f930c02d4 100644 --- a/internal/handler/searchbot_test.go +++ b/internal/handler/searchbot_test.go @@ -28,6 +28,7 @@ import ( "ragflow/internal/common" "ragflow/internal/entity" + "ragflow/internal/permission" "ragflow/internal/service" "github.com/gin-gonic/gin" @@ -36,16 +37,14 @@ import ( // mockChunkService implements service.Retriever for testing. // It captures the last request received so tests can verify field mapping. type mockChunkService struct { - retrievalTestFn func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) + retrievalTestFn func(req *service.RetrievalTestRequest) (*service.RetrievalTestResponse, error) LastReq *service.RetrievalTestRequest - LastUserID string } -func (m *mockChunkService) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { +func (m *mockChunkService) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest) (*service.RetrievalTestResponse, error) { m.LastReq = req - m.LastUserID = userID if m.retrievalTestFn != nil { - return m.retrievalTestFn(req, userID) + return m.retrievalTestFn(req) } return &service.RetrievalTestResponse{ Chunks: []map[string]interface{}{{"docnm_kwd": "test", "content_with_weight": "content"}}, @@ -98,9 +97,6 @@ func TestSearchBotsRetrieval_Basic(t *testing.T) { if mockSvc.LastReq.Question != "test question" { t.Errorf("Question = %q, want \"test question\"", mockSvc.LastReq.Question) } - if mockSvc.LastUserID != "user1" { - t.Errorf("userID = %q, want \"user1\"", mockSvc.LastUserID) - } } func TestSearchBotsRetrieval_MissingKbID(t *testing.T) { @@ -167,7 +163,7 @@ func TestSearchBotsRetrieval_NoAuth(t *testing.T) { func TestSearchBotsRetrieval_ServiceError(t *testing.T) { h, _, r := setupSearchbotsTest("user1") h.chunkSvc = &mockChunkService{ - retrievalTestFn: func(req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { + retrievalTestFn: func(req *service.RetrievalTestRequest) (*service.RetrievalTestResponse, error) { return nil, errors.New("db error") }, } @@ -467,7 +463,7 @@ type fakeChunkRetriever struct { err error } -func (f *fakeChunkRetriever) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { +func (f *fakeChunkRetriever) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest) (*service.RetrievalTestResponse, error) { if f.err != nil { return nil, f.err } @@ -538,6 +534,7 @@ func TestAskHandler_NoChatModel(t *testing.T) { strings.NewReader(`{"question": "test", "kb_ids": ["kb1"]}`)) c.Request.Header.Set("Content-Type", "application/json") c.Set("user", &entity.User{ID: "user-1"}) + c.Request = c.Request.WithContext(permission.WithSubject(c.Request.Context(), permission.Subject{UserID: "user-1"})) h.Ask(c) body := buf.String() if !strings.Contains(body, "chat model not configured") { diff --git a/internal/handler/tenant.go b/internal/handler/tenant.go index 61432b6499..31b420b19c 100644 --- a/internal/handler/tenant.go +++ b/internal/handler/tenant.go @@ -39,17 +39,17 @@ import ( // TenantHandler tenant handler type TenantHandler struct { - tenantService *service.TenantService - userService *service.UserService - datasetService *dataset.DatasetService + tenantService *service.TenantService + userTenantService *service.UserTenantService + datasetService *dataset.DatasetService } // NewTenantHandler create tenant handler -func NewTenantHandler(tenantService *service.TenantService, userService *service.UserService, datasetService *dataset.DatasetService) *TenantHandler { +func NewTenantHandler(tenantService *service.TenantService, userTenantService *service.UserTenantService, datasetService *dataset.DatasetService) *TenantHandler { return &TenantHandler{ - tenantService: tenantService, - userService: userService, - datasetService: datasetService, + tenantService: tenantService, + userTenantService: userTenantService, + datasetService: datasetService, } } @@ -84,9 +84,17 @@ func (h *TenantHandler) setDefaultModels(c *gin.Context, wrapModels bool) { } ctx := c.Request.Context() - err := h.tenantService.SetTenantDefaultModels(ctx, user.ID, req.ModelProvider, req.ModelInstance, req.ModelName, req.ModelType, req.ModelID) + code, err := h.userTenantService.SetModelSettings(ctx, user.ID, &service.SetModelSettingsRequest{ + DefaultModel: &service.DefaultModelSelection{ + ModelProvider: req.ModelProvider, + ModelInstance: req.ModelInstance, + ModelName: req.ModelName, + ModelType: req.ModelType, + ModelID: req.ModelID, + }, + }) if err != nil { - common.ResponseWithCodeData(c, common.CodeExceptionError, false, err.Error()) + common.ResponseWithCodeData(c, code, false, err.Error()) return } @@ -98,7 +106,7 @@ func (h *TenantHandler) setDefaultModels(c *gin.Context, wrapModels bool) { common.SuccessNoData(c, "success") } -// GetDefaultModels returns the tenant's default model selections. The +// GetDefaultModels returns the user's model selections for the current tenant. The // response wraps the model list under `data.models` to mirror the // Python `list_tenant_default_models` contract (api/apps/restful_apis/ // models_api.py:84). The frontend hook `useFetchDefaultModels` @@ -111,7 +119,7 @@ func (h *TenantHandler) GetDefaultModels(c *gin.Context) { } ctx := c.Request.Context() - defaultModels, err := h.tenantService.ListTenantDefaultModels(ctx, user.ID) + defaultModels, err := h.userTenantService.ListDefaultModels(ctx, user.ID) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, false, err.Error()) return @@ -142,7 +150,7 @@ func (h *TenantHandler) TenantInfo(c *gin.Context) { } ctx := c.Request.Context() - tenantInfo, err := h.tenantService.GetTenantInfo(ctx, user.ID) + tenantInfo, err := h.userTenantService.GetCurrentTenantInfo(ctx, user.ID) if err != nil { common.ResponseWithCodeData(c, common.CodeExceptionError, false, err.Error()) return diff --git a/internal/handler/user.go b/internal/handler/user.go index c45874a59a..a163ac270b 100644 --- a/internal/handler/user.go +++ b/internal/handler/user.go @@ -35,13 +35,15 @@ import ( // UserHandler user handler type UserHandler struct { - userService *service.UserService + userService *service.UserService + userTenantService *service.UserTenantService } // NewUserHandler create user handler -func NewUserHandler(userService *service.UserService) *UserHandler { +func NewUserHandler(userService *service.UserService, userTenantService *service.UserTenantService) *UserHandler { return &UserHandler{ - userService: userService, + userService: userService, + userTenantService: userTenantService, } } @@ -473,7 +475,7 @@ func (h *UserHandler) ChangePassword(c *gin.Context) { // @Description Update tenant model configuration // @Tags users // @Security ApiKeyAuth -// @Param request body service.SetTenantInfoRequest true "tenant info" +// @Param request body service.SetModelSettingsRequest true "tenant model settings" // @Success 200 {object} map[string]interface{} // @Router /v1/user/set_tenant_info [post] func (h *UserHandler) SetTenantInfo(c *gin.Context) { @@ -524,7 +526,7 @@ func (h *UserHandler) SetTenantInfo(c *gin.Context) { return } - req := service.SetTenantInfoRequest{Raw: payload} + req := service.SetModelSettingsRequest{Raw: payload} if value, ok := payload["tenant_id"].(string); ok { req.TenantID = &value } @@ -547,7 +549,7 @@ func (h *UserHandler) SetTenantInfo(c *gin.Context) { req.TTSID = &value } - code, err := h.userService.SetTenantInfo(ctx, user.ID, &req) + code, err := h.userTenantService.SetModelSettings(ctx, user.ID, &req) if err != nil { common.ResponseWithCodeData(c, code, nil, err.Error()) operationLog.Code = int(code) diff --git a/internal/ingestion/component/dispatch_model.go b/internal/ingestion/component/dispatch_model.go index 8c37cd885f..e433de0f17 100644 --- a/internal/ingestion/component/dispatch_model.go +++ b/internal/ingestion/component/dispatch_model.go @@ -31,13 +31,14 @@ import ( "ragflow/internal/entity" modelModule "ragflow/internal/entity/models" "ragflow/internal/ingestion/component/schema" + "ragflow/internal/modelcontext" "ragflow/internal/permission" "gorm.io/gorm" ) type tenantModelExtra struct { - MaxTokens *int `json:"max_tokens"` + MaxTokens json.RawMessage `json:"max_tokens"` } type resolvedModelConfig struct { @@ -114,8 +115,7 @@ var resolveTenantOCRModelByProvider = defaultResolveTenantOCRModelByProvider // model is resolved by ID; its resolution error is returned rather than trying // another model. func defaultResolveTenantOCRModelByProvider(ctx context.Context, db *gorm.DB, tenantID string, providerName string) (modelModule.ModelDriver, string, *modelModule.APIConfig, int, error) { - providerDAO := dao.NewTenantModelProviderDAO() - provider, err := providerDAO.GetByTenantIDAndProviderName(ctx, db, tenantID, providerName) + provider, err := getTenantModelProviderByName(ctx, db, tenantID, providerName) if err != nil { // Some OCR capabilities are registered under sibling provider names // (cloud vs. local). Tolerate the alternate spelling before giving up @@ -124,7 +124,7 @@ func defaultResolveTenantOCRModelByProvider(ctx context.Context, db *gorm.DB, te if alias == providerName { continue } - if p, aerr := providerDAO.GetByTenantIDAndProviderName(ctx, db, tenantID, alias); aerr == nil { + if p, aerr := getTenantModelProviderByName(ctx, db, tenantID, alias); aerr == nil { provider, err = p, nil break } @@ -240,8 +240,12 @@ func resolveModelConfigByID(ctx context.Context, db *gorm.DB, tenantID string, m // Scope the lookup itself so missing and foreign models are indistinguishable // before exposing status, type, instance configuration, or driver details. + owner, err := dao.NewUserTenantDAO().GetOwnerByTenantID(ctx, db, tenantID) + if err != nil { + return nil, fmt.Errorf("resolve tenant owner for model lookup: %w", err) + } providers := db.Session(&gorm.Session{}).Model(&entity.TenantModelProvider{}). - Select("id").Where("tenant_id = ?", tenantID) + Select("id").Where("tenant_id = ? AND user_id = ?", tenantID, owner.UserID) scopedDB := db.Session(&gorm.Session{}).Where("provider_id IN (?)", providers) modelObj, err := modelDAO.GetByID(ctx, scopedDB, modelID) if err != nil { @@ -305,8 +309,8 @@ func resolveModelConfigByID(ctx context.Context, db *gorm.DB, tenantID string, m if err := json.Unmarshal([]byte(modelObj.Extra), &tenantExtra); err != nil { return nil, err } - if tenantExtra.MaxTokens != nil && *tenantExtra.MaxTokens > 0 { - contextLength = *tenantExtra.MaxTokens + if maxTokens := modelcontext.MaxTokensFromJSON(tenantExtra.MaxTokens); maxTokens > 0 { + contextLength = maxTokens } } apiConfig := &modelModule.APIConfig{ApiKey: &apiKey, Region: ®ion, BaseURL: &baseURL} @@ -323,11 +327,10 @@ func resolveModelConfigFromProviderInstance(ctx context.Context, db *gorm.DB, te } catalogName := modelModule.CatalogProviderName(providerName) - providerDAO := dao.NewTenantModelProviderDAO() instanceDAO := dao.NewTenantModelInstanceDAO() modelDAO := dao.NewTenantModelDAO() - provider, err := providerDAO.GetByTenantIDAndProviderName(ctx, db, tenantID, providerName) + provider, err := getTenantModelProviderByName(ctx, db, tenantID, providerName) if err != nil { return nil, "", nil, 0, fmt.Errorf("provider %q lookup failed: %w", providerName, err) } @@ -378,8 +381,8 @@ func resolveModelConfigFromProviderInstance(ctx context.Context, db *gorm.DB, te if err := json.Unmarshal([]byte(modelObj.Extra), &tenantExtra); err != nil { return nil, "", nil, 0, err } - if tenantExtra.MaxTokens != nil && *tenantExtra.MaxTokens > 0 { - contextLength = *tenantExtra.MaxTokens + if maxTokens := modelcontext.MaxTokensFromJSON(tenantExtra.MaxTokens); maxTokens > 0 { + contextLength = maxTokens } } apiConfig := &modelModule.APIConfig{ApiKey: &apiKey, Region: ®ion, BaseURL: &baseURL} @@ -437,3 +440,11 @@ func newModelDriverForBaseURLLocal(driver modelModule.ModelDriver, providerName, func errorsIsRecordNotFound(err error) bool { return err != nil && (errors.Is(err, gorm.ErrRecordNotFound) || strings.Contains(err.Error(), gorm.ErrRecordNotFound.Error())) } + +func getTenantModelProviderByName(ctx context.Context, db *gorm.DB, tenantID, providerName string) (*entity.TenantModelProvider, error) { + owner, err := dao.NewUserTenantDAO().GetOwnerByTenantID(ctx, db, tenantID) + if err != nil { + return nil, err + } + return dao.NewTenantModelProviderDAO().GetByTenantIDAndUserIDAndProviderName(ctx, db, tenantID, owner.UserID, providerName) +} diff --git a/internal/ingestion/component/dispatch_model_test.go b/internal/ingestion/component/dispatch_model_test.go index 25fb4f9ce1..bb838cd8b3 100644 --- a/internal/ingestion/component/dispatch_model_test.go +++ b/internal/ingestion/component/dispatch_model_test.go @@ -66,12 +66,12 @@ func TestResolveModelConfigForStoredMistralOCRProvider(t *testing.T) { if err != nil { t.Fatalf("open sqlite: %v", err) } - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatalf("migrate: %v", err) } const storedKey = `{"MISTRAL_OCR_API_KEY": "sk-env"}` for _, row := range []interface{}{ - &entity.TenantModelProvider{ID: "provider-mistral-ocr", TenantID: "tenant-1", ProviderName: "Mistral OCR"}, + &entity.TenantModelProvider{ID: "provider-mistral-ocr", UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "Mistral OCR"}, &entity.TenantModelInstance{ID: "instance-mistral-ocr", ProviderID: "provider-mistral-ocr", InstanceName: "default", APIKey: storedKey, Status: "active", Extra: "{}"}, &entity.TenantModel{ID: "model-mistral-ocr", ProviderID: "provider-mistral-ocr", InstanceID: "instance-mistral-ocr", ModelName: "mistral-ocr-latest", ModelType: int(entity.ModelTypeOCR), Status: "active"}, } { @@ -262,12 +262,12 @@ func TestResolveModelConfigByIDScopedQuerySQL(t *testing.T) { t.Fatalf("model queries=%d, want one scoped lookup: %+v", len(queries), queries) } query := queries[0] - const wantSQL = "SELECT * FROM `tenant_model` WHERE provider_id IN (SELECT `id` FROM `tenant_model_provider` WHERE tenant_id = ?) AND id = ? LIMIT 1" + const wantSQL = "SELECT * FROM `tenant_model` WHERE provider_id IN (SELECT `id` FROM `tenant_model_provider` WHERE tenant_id = ? AND user_id = ?) AND id = ? LIMIT 1" if query.sql != wantSQL { t.Errorf("scoped SQL=%q, want %q", query.sql, wantSQL) } - if len(query.vars) != 2 || query.vars[0] != tenantID || query.vars[1] != modelID { - t.Errorf("bound vars=%#v, want tenant ID then model ID", query.vars) + if len(query.vars) != 3 || query.vars[0] != tenantID || query.vars[2] != modelID { + t.Errorf("bound vars=%#v, want tenant ID, owner user ID, then model ID", query.vars) } if strings.Contains(query.sql, tenantID) || strings.Contains(query.sql, modelID) { t.Errorf("an ID was interpolated into SQL: %q", query.sql) @@ -353,7 +353,7 @@ func openDispatchModelTestDB(t *testing.T) *gorm.DB { if err != nil { t.Fatalf("open sqlite: %v", err) } - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatalf("migrate: %v", err) } return db @@ -367,7 +367,7 @@ func openDispatchModelTestDB(t *testing.T) *gorm.DB { func TestResolveModelConfig_InactiveInstanceIsAnExplicitFailure(t *testing.T) { db := openDispatchModelTestDB(t) for _, row := range []interface{}{ - &entity.TenantModelProvider{ID: "provider-openai", TenantID: "tenant-1", ProviderName: "OpenAI"}, + &entity.TenantModelProvider{ID: "provider-openai", UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "OpenAI"}, &entity.TenantModelInstance{ID: "instance-default", ProviderID: "provider-openai", InstanceName: "default", APIKey: `{"OPENAI_API_KEY": "sk-inactive"}`, Status: "inactive"}, &entity.TenantModel{ID: "model-1", ProviderID: "provider-openai", InstanceID: "instance-default", ModelName: "gpt-4o", ModelType: int(entity.ModelTypeChat), Status: "active"}, } { @@ -403,7 +403,7 @@ func TestResolveModelConfig_InactiveInstanceIsAnExplicitFailure(t *testing.T) { func TestResolveModelConfig_MaxTokensOverrideSetsContextWindowOnly(t *testing.T) { db := openDispatchModelTestDB(t) for _, row := range []interface{}{ - &entity.TenantModelProvider{ID: "provider-openai", TenantID: "tenant-1", ProviderName: "OpenAI"}, + &entity.TenantModelProvider{ID: "provider-openai", UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "OpenAI"}, &entity.TenantModelInstance{ID: "instance-default", ProviderID: "provider-openai", InstanceName: "default", Status: "active"}, &entity.TenantModel{ID: "model-1", ProviderID: "provider-openai", InstanceID: "instance-default", ModelName: "gpt-4o", ModelType: int(entity.ModelTypeChat), Status: "active"}, } { diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 7296a51afb..8e895bda59 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -55,12 +55,12 @@ import ( "ragflow/internal/agent/runtime" "ragflow/internal/common" - "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/entity/models" "ragflow/internal/ingestion/chunkcache" "ragflow/internal/ingestion/component/progress" "ragflow/internal/ingestion/component/schema" + "ragflow/internal/modelcontext" "ragflow/internal/tokenizer" "ragflow/internal/utility" ) @@ -1669,7 +1669,7 @@ func extractorContextLength(ctx context.Context, db *gorm.DB, llmID string) int if llmID == "" { return 0 } - return dao.ResolveModelContentLength(ctx, db, tid, llmID, "", "") + return modelcontext.ResolveModelContentLength(ctx, db, tid, llmID, "", "") } // defaultChatModelRef returns the tenant's default chat model reference — diff --git a/internal/ingestion/component/extractor_test.go b/internal/ingestion/component/extractor_test.go index ea5d441190..bd9d05434b 100644 --- a/internal/ingestion/component/extractor_test.go +++ b/internal/ingestion/component/extractor_test.go @@ -1390,6 +1390,7 @@ func seedExtractorContextModel(t *testing.T, db *gorm.DB, tenantLLMID string) { } if err := db.Create(&entity.TenantModelProvider{ ID: "provider-openai", + UserID: "tenant-1", ProviderName: "OpenAI", TenantID: "tenant-1", }).Error; err != nil { diff --git a/internal/ingestion/component/image_ocr_test.go b/internal/ingestion/component/image_ocr_test.go index bfee431cb3..d46b599a24 100644 --- a/internal/ingestion/component/image_ocr_test.go +++ b/internal/ingestion/component/image_ocr_test.go @@ -146,10 +146,10 @@ func imageOCRFixture(t *testing.T, provider string, fail bool) (*gorm.DB, *image } sqlDB.SetMaxOpenConns(1) t.Cleanup(func() { _ = sqlDB.Close() }) - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatal(err) } - rows := []any{&entity.TenantModelProvider{ID: "provider", TenantID: "tenant-1", ProviderName: provider}} + rows := []any{&entity.TenantModelProvider{ID: "provider", UserID: "tenant-1", TenantID: "tenant-1", ProviderName: provider}} for _, instance := range []struct{ id, url string }{{"a", decoy.URL}, {"b", server.URL}} { rows = append(rows, &entity.TenantModelInstance{ID: "instance-" + instance.id, ProviderID: "provider", InstanceName: instance.id, diff --git a/internal/ingestion/component/parser_dispatch_test.go b/internal/ingestion/component/parser_dispatch_test.go index ee0590576d..4746ea959d 100644 --- a/internal/ingestion/component/parser_dispatch_test.go +++ b/internal/ingestion/component/parser_dispatch_test.go @@ -1592,12 +1592,12 @@ func setupMinerUVisionDispatchDB(t *testing.T, baseURL string) (*gorm.DB, string if err != nil { t.Fatalf("open sqlite: %v", err) } - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatalf("migrate model tables: %v", err) } modelID := "061436bbd8474d54be5bba3efbeac109" rows := []any{ - &entity.TenantModelProvider{ID: "provider-mineru-1", TenantID: "test-tenant", ProviderName: "MinerU"}, + &entity.TenantModelProvider{ID: "provider-mineru-1", UserID: "test-tenant", TenantID: "test-tenant", ProviderName: "MinerU"}, &entity.TenantModelInstance{ID: "instance-mineru-1", ProviderID: "provider-mineru-1", InstanceName: "default", APIKey: "mineru-secret", Status: "active", Extra: `{"base_url":"` + baseURL + `"}`}, &entity.TenantModel{ID: modelID, ProviderID: "provider-mineru-1", InstanceID: "instance-mineru-1", ModelName: "MinerU-model", ModelType: int(entity.ModelTypeOCR), Status: "active"}, } @@ -1766,11 +1766,11 @@ func TestResolveTenantOCRModelByProvider_MinerUProviderAlias(t *testing.T) { if err != nil { t.Fatalf("open sqlite: %v", err) } - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatalf("migrate model tables: %v", err) } rows := []any{ - &entity.TenantModelProvider{ID: "provider-minerunet-1", TenantID: "test-tenant", ProviderName: "MinerU.Net"}, + &entity.TenantModelProvider{ID: "provider-minerunet-1", UserID: "test-tenant", TenantID: "test-tenant", ProviderName: "MinerU.Net"}, &entity.TenantModelInstance{ID: "instance-minerunet-1", ProviderID: "provider-minerunet-1", InstanceName: "default", APIKey: "mineru-secret", Status: "active", Extra: "{}"}, &entity.TenantModel{ID: "model-minerunet-1", ProviderID: "provider-minerunet-1", InstanceID: "instance-minerunet-1", ModelName: "MinerU-model", ModelType: int(entity.ModelTypeOCR), Status: "active"}, } diff --git a/internal/ingestion/component/pdf_parser_model.go b/internal/ingestion/component/pdf_parser_model.go index beaed4c973..6906aa3316 100644 --- a/internal/ingestion/component/pdf_parser_model.go +++ b/internal/ingestion/component/pdf_parser_model.go @@ -86,7 +86,12 @@ func maybeDispatchTenantPDFParser(ctx context.Context, db *gorm.DB, filename str var instance *entity.TenantModelInstance var err error if composite { - provider, err = providerDAO.GetByTenantIDAndProviderName(ctx, db, tenantID, providerName) + owner, ownerErr := dao.NewUserTenantDAO().GetOwnerByTenantID(ctx, db, tenantID) + if ownerErr != nil { + return parser.ParseResult{}, true, fmt.Errorf("parser: tenant owner lookup: %w", ownerErr) + } + ownerUserID := owner.UserID + provider, err = providerDAO.GetByTenantIDAndUserIDAndProviderName(ctx, db, tenantID, ownerUserID, providerName) if err == nil { instance, err = instanceDAO.GetByProviderIDAndInstanceName(ctx, db, provider.ID, instanceName) } diff --git a/internal/ingestion/component/pdf_parser_model_test.go b/internal/ingestion/component/pdf_parser_model_test.go index 0fa2cc7e45..65d0f5508f 100644 --- a/internal/ingestion/component/pdf_parser_model_test.go +++ b/internal/ingestion/component/pdf_parser_model_test.go @@ -98,10 +98,10 @@ func tenantPDFParserFixture(t *testing.T, provider string) (*gorm.DB, string, ch } sqlDB.SetMaxOpenConns(1) t.Cleanup(func() { _ = sqlDB.Close() }) - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatal(err) } - if err := db.Create(&entity.TenantModelProvider{ID: "provider", ProviderName: provider, TenantID: "tenant-1"}).Error; err != nil { + if err := db.Create(&entity.TenantModelProvider{ID: "provider", ProviderName: provider, UserID: "tenant-1", TenantID: "tenant-1"}).Error; err != nil { t.Fatal(err) } for _, inst := range []struct{ id, url string }{{"a", decoy.URL}, {"b", server.URL}} { diff --git a/internal/ingestion/pipeline/image_ocr_test.go b/internal/ingestion/pipeline/image_ocr_test.go index 8b4646e65c..80b96ea3ed 100644 --- a/internal/ingestion/pipeline/image_ocr_test.go +++ b/internal/ingestion/pipeline/image_ocr_test.go @@ -98,7 +98,7 @@ func TestImageOCRSavedConfigToSearchableChunk(t *testing.T) { } sqlDB.SetMaxOpenConns(1) defer sqlDB.Close() - if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatal(err) } registry, err := DefaultRegistry() @@ -120,7 +120,7 @@ func TestImageOCRSavedConfigToSearchableChunk(t *testing.T) { }}) for _, row := range []any{ &entity.Knowledgebase{ID: "kb-1", TenantID: "tenant-1", Name: "Image OCR fixture", ParserConfig: config}, - &entity.TenantModelProvider{ID: "ocr-provider", TenantID: "tenant-1", ProviderName: "PaddleOCR.local"}, + &entity.TenantModelProvider{ID: "ocr-provider", UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "PaddleOCR.local"}, &entity.TenantModelInstance{ID: "ocr-instance", ProviderID: "ocr-provider", InstanceName: "selected", Status: "active", APIKey: "fixture-key", Extra: fmt.Sprintf(`{"base_url":%q}`, server.URL)}, &entity.TenantModel{ID: "ocr-model", ProviderID: "ocr-provider", InstanceID: "ocr-instance", diff --git a/internal/ingestion/task/embedder.go b/internal/ingestion/task/embedder.go index 8c5ff04a95..5f86a3ed7b 100644 --- a/internal/ingestion/task/embedder.go +++ b/internal/ingestion/task/embedder.go @@ -707,7 +707,7 @@ func init() { return kb.EmbdID, nil }, func(ctx context.Context, tenantID, embdID string) (*models.EmbeddingModel, error) { - embeddingModel, err := modelFactory.NewEmbeddingModel(ctx, service.ModelAccess{TenantID: tenantID}, embdID) + embeddingModel, err := modelFactory.NewEmbeddingModel(ctx, service.Identity{TenantID: tenantID}, embdID) if err != nil { return nil, err } diff --git a/internal/ingestion/task/knowledge_compiler_wiring.go b/internal/ingestion/task/knowledge_compiler_wiring.go index 281cd92660..0cec413340 100644 --- a/internal/ingestion/task/knowledge_compiler_wiring.go +++ b/internal/ingestion/task/knowledge_compiler_wiring.go @@ -582,7 +582,7 @@ func newKnowledgeCompilerDepsResolver() kc.DepsResolver { // actually run instead of failing / falling back to a no-op. Use the // tenant model ID so later model resolution can use the same persisted // tenant model directly. - defaultInfo, derr := modelFactory.ResolveDefaultInfo(modelCtx, service.ModelAccess{TenantID: tenantID}, entity.ModelTypeChat) + defaultInfo, derr := modelFactory.ResolveDefaultInfo(modelCtx, service.Identity{TenantID: tenantID}, entity.ModelTypeChat) if derr != nil || defaultInfo == nil || strings.TrimSpace(defaultInfo.ID) == "" { return kc.Deps{}, fmt.Errorf("knowledge_compiler: llm_id is empty and no tenant default chat model available: %w", derr) } @@ -596,7 +596,7 @@ func newKnowledgeCompilerDepsResolver() kc.DepsResolver { // generation cap and must not be used as the input budget. llmMax := kc.DefaultLLMContextLength // Bound the model-config lookup and inherit cancellation from the task. - modelInfo, contextLengthErr := modelFactory.ResolveInfo(modelCtx, service.ModelAccess{TenantID: tenantID}, entity.ModelTypeChat, llmID) + modelInfo, contextLengthErr := modelFactory.ResolveInfo(modelCtx, service.Identity{TenantID: tenantID}, entity.ModelTypeChat, llmID) if contextLengthErr == nil && modelInfo != nil && modelInfo.ContextLength > 0 { llmMax = modelInfo.ContextLength } @@ -642,7 +642,7 @@ func (c *kcChatInvoker) Chat(ctx context.Context, req kc.ChatRequest) (*kc.ChatR if req.LLMID != "" { llmID = req.LLMID } - chatModel, err := c.factory.NewChatModel(ctx, service.ModelAccess{TenantID: c.tenantID}, llmID) + chatModel, err := c.factory.NewChatModel(ctx, service.Identity{TenantID: c.tenantID}, llmID) if err != nil { return nil, err } @@ -822,7 +822,7 @@ func (e *kcEmbedder) Encode(ctx context.Context, texts []string) ([][]float32, e // loudly instead of silently producing empty vectors. func (e *kcEmbedder) resolveModel(ctx context.Context) (*models.EmbeddingModel, error) { if embdID := strings.TrimSpace(e.embdID); embdID != "" { - embeddingModel, err := e.factory.NewEmbeddingModel(ctx, service.ModelAccess{TenantID: e.tenantID}, embdID) + embeddingModel, err := e.factory.NewEmbeddingModel(ctx, service.Identity{TenantID: e.tenantID}, embdID) if err != nil { return nil, fmt.Errorf("knowledge_compiler: resolve embedding model: %w", err) } @@ -831,7 +831,7 @@ func (e *kcEmbedder) resolveModel(ctx context.Context) (*models.EmbeddingModel, } return embeddingModel, nil } - embeddingModel, err := e.factory.NewDefaultEmbeddingModel(ctx, service.ModelAccess{TenantID: e.tenantID}) + embeddingModel, err := e.factory.NewDefaultEmbeddingModel(ctx, service.Identity{TenantID: e.tenantID}) if err != nil { return nil, fmt.Errorf("knowledge_compiler: embedding_model is required and no tenant default embedding model is set: %w", err) } diff --git a/internal/dao/model_context.go b/internal/modelcontext/context_length.go similarity index 73% rename from internal/dao/model_context.go rename to internal/modelcontext/context_length.go index 45be63c6dc..152bbf775b 100644 --- a/internal/dao/model_context.go +++ b/internal/modelcontext/context_length.go @@ -14,7 +14,7 @@ // limitations under the License. // -package dao +package modelcontext import ( "context" @@ -24,13 +24,14 @@ import ( "gorm.io/gorm" + "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/entity/models" ) // ResolveModelContentLength returns the chat model's effective context window -// (context_length) in tokens for modelRef — a tenant_model UUID or a -// composite "model@provider" / "model@instance@provider" reference — or 0 +// (context_length) in tokens for modelRef - a tenant_model UUID or a +// composite "model@provider" / "model@instance@provider" reference - or 0 // when it cannot be resolved. tenantID scopes the composite-reference lookup // to the tenant's own provider/instance/model rows. driver and modelName are // an optional provider catalog fallback used when modelRef is not a @@ -45,12 +46,11 @@ import ( // context_length (via the tenant row, the composite reference parts, or // the resolved driver + model name). // -// This is the shared implementation behind the agent LLM component, the -// ingestion Extractor, and service.ModelProviderService — those packages -// cannot import each other (import cycles), so the lookup lives here. +// This is the shared implementation behind the Agent LLM component and the +// Ingestion Extractor, which cannot depend on the Service package. func ResolveModelContentLength(ctx context.Context, db *gorm.DB, tenantID, modelRef, driver, modelName string) int { if db == nil { - db = DB + db = dao.DB } pureName, instanceName, providerName, composite := splitCompositeModelRef(modelRef) @@ -64,31 +64,31 @@ func ResolveModelContentLength(ctx context.Context, db *gorm.DB, tenantID, model } // 2. Catalog fallbacks (only when there is no custom override). - // 2a. Active tenant row → its provider row → catalog context_length. + // 2a. Active tenant row, its provider row, then catalog context_length. if obj != nil && obj.Status == "active" { - if provider, err := NewTenantModelProviderDAO().GetByID(ctx, db, obj.ProviderID); err == nil && provider != nil { - if mdl, err := GetModelProviderManager().GetModelByName(models.CatalogProviderName(provider.ProviderName), obj.ModelName); err == nil && mdl.ContextLength != nil { + if provider, err := dao.NewTenantModelProviderDAO().GetByID(ctx, db, obj.ProviderID); err == nil && provider != nil { + if mdl, err := dao.GetModelProviderManager().GetModelByName(models.CatalogProviderName(provider.ProviderName), obj.ModelName); err == nil && mdl.ContextLength != nil { return *mdl.ContextLength } } } - // 2b. Composite reference with no tenant row → catalog by reference parts. + // 2b. Composite reference with no tenant row; query catalog by reference parts. if composite { - if mdl, err := GetModelProviderManager().GetModelByName(models.CatalogProviderName(providerName), pureName); err == nil && mdl.ContextLength != nil { + if mdl, err := dao.GetModelProviderManager().GetModelByName(models.CatalogProviderName(providerName), pureName); err == nil && mdl.ContextLength != nil { return *mdl.ContextLength } } // 2c. Resolved driver + bare model name: fallback when modelRef is a // tenant id that could not be resolved without a database. if driver != "" && modelName != "" { - if mdl, err := GetModelProviderManager().GetModelByName(models.CatalogProviderName(driver), modelName); err == nil && mdl.ContextLength != nil { + if mdl, err := dao.GetModelProviderManager().GetModelByName(models.CatalogProviderName(driver), modelName); err == nil && mdl.ContextLength != nil { return *mdl.ContextLength } } return 0 } -// lookupTenantModel resolves the tenant_model row for modelRef — by UUID, or +// lookupTenantModel resolves the tenant_model row for modelRef by UUID, or // for a composite reference through the tenant's provider/instance rows (chat // model type). Returns nil when the row cannot be resolved. func lookupTenantModel(ctx context.Context, db *gorm.DB, tenantID, modelRef, pureName, instanceName, providerName string, composite bool) *entity.TenantModel { @@ -99,15 +99,20 @@ func lookupTenantModel(ctx context.Context, db *gorm.DB, tenantID, modelRef, pur if tenantID == "" { return nil } - provider, err := NewTenantModelProviderDAO().GetByTenantIDAndProviderName(ctx, db, tenantID, providerName) + owner, err := dao.NewUserTenantDAO().GetOwnerByTenantID(ctx, db, tenantID) + if err != nil { + return nil + } + ownerUserID := owner.UserID + provider, err := dao.NewTenantModelProviderDAO().GetByTenantIDAndUserIDAndProviderName(ctx, db, tenantID, ownerUserID, providerName) if err != nil || provider == nil { return nil } - instance, err := NewTenantModelInstanceDAO().GetByProviderIDAndInstanceName(ctx, db, provider.ID, instanceName) + instance, err := dao.NewTenantModelInstanceDAO().GetByProviderIDAndInstanceName(ctx, db, provider.ID, instanceName) if err != nil || instance == nil || instance.Status != "active" { return nil } - obj, err := NewTenantModelDAO().GetByProviderIDAndInstanceIDAndModelTypeAndModelName(ctx, db, provider.ID, instance.ID, int(entity.ModelTypeChat), pureName) + obj, err := dao.NewTenantModelDAO().GetByProviderIDAndInstanceIDAndModelTypeAndModelName(ctx, db, provider.ID, instance.ID, int(entity.ModelTypeChat), pureName) if err != nil || obj == nil { return nil } @@ -116,7 +121,7 @@ func lookupTenantModel(ctx context.Context, db *gorm.DB, tenantID, modelRef, pur if modelRef == "" { return nil } - obj, err := NewTenantModelDAO().GetByID(ctx, db, modelRef) + obj, err := dao.NewTenantModelDAO().GetByID(ctx, db, modelRef) if err != nil || obj == nil { return nil } @@ -130,7 +135,7 @@ func lookupTenantModel(ctx context.Context, db *gorm.DB, tenantID, modelRef, pur // modelExtraMaxTokens returns the per-model "max_tokens" override from the // tenant_model.extra JSON, or 0 when absent/invalid/non-positive. Mirrors -// Python's model_extra.get("max_tokens") — the custom context-window length a +// Python's model_extra.get("max_tokens"): the custom context-window length a // user configures for a model. Both JSON numbers and numeric strings are // accepted (some callers persist extra.max_tokens as a string). func modelExtraMaxTokens(extra string) int { @@ -143,17 +148,25 @@ func modelExtraMaxTokens(extra string) int { if err := json.Unmarshal([]byte(extra), &m); err != nil || len(m.MaxTokens) == 0 { return 0 } + return MaxTokensFromJSON(m.MaxTokens) +} + +// MaxTokensFromJSON parses a positive max_tokens JSON number or numeric string. +func MaxTokensFromJSON(raw json.RawMessage) int { + if len(raw) == 0 { + return 0 + } // Accept JSON numbers (int or float, e.g. 32000 / 32000.0) and numeric // strings ("32000"); json.Number keeps the raw lexical form. var num json.Number - if err := json.Unmarshal(m.MaxTokens, &num); err == nil { + if err := json.Unmarshal(raw, &num); err == nil { if f, err := strconv.ParseFloat(num.String(), 64); err == nil && f > 0 { return int(f) } return 0 } var str string - if err := json.Unmarshal(m.MaxTokens, &str); err == nil { + if err := json.Unmarshal(raw, &str); err == nil { if n, err := strconv.Atoi(strings.TrimSpace(str)); err == nil && n > 0 { return n } diff --git a/internal/dao/model_context_test.go b/internal/modelcontext/context_length_test.go similarity index 67% rename from internal/dao/model_context_test.go rename to internal/modelcontext/context_length_test.go index be304e0d3f..03d8fa7c90 100644 --- a/internal/dao/model_context_test.go +++ b/internal/modelcontext/context_length_test.go @@ -14,22 +14,52 @@ // limitations under the License. // -package dao +package modelcontext import ( + "context" + "encoding/json" "testing" "github.com/glebarez/sqlite" "gorm.io/gorm" + "ragflow/internal/common" "ragflow/internal/entity" ) +func TestMaxTokensFromJSON(t *testing.T) { + tests := []struct { + name string + json json.RawMessage + want int + }{ + {name: "integer", json: json.RawMessage(`32000`), want: 32000}, + {name: "fractional number", json: json.RawMessage(`32000.9`), want: 32000}, + {name: "numeric string", json: json.RawMessage(`" 32000 "`), want: 32000}, + {name: "zero", json: json.RawMessage(`0`)}, + {name: "negative", json: json.RawMessage(`-1`)}, + {name: "invalid string", json: json.RawMessage(`"large"`)}, + {name: "null", json: json.RawMessage(`null`)}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if got := MaxTokensFromJSON(test.json); got != test.want { + t.Fatalf("MaxTokensFromJSON(%s) = %d, want %d", test.json, got, test.want) + } + }) + } +} + +func resolveTestContextLength(ctx context.Context, db *gorm.DB, tenantID, modelRef, driver, modelName string) int { + return ResolveModelContentLength(ctx, db, tenantID, modelRef, driver, modelName) +} + // TestResolveModelContentLength_CompositeReference resolves context_length // for a composite "model@provider" reference from the provider catalog. func TestResolveModelContentLength_CompositeReference(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "", "gpt-4o@openai", "", ""); got != 128000 { - t.Fatalf("ResolveModelContentLength(gpt-4o@openai) = %d, want 128000", got) + if got := resolveTestContextLength(t.Context(), nil, "", "gpt-4o@openai", "", ""); got != 128000 { + t.Fatalf("resolveTestContextLength(gpt-4o@openai) = %d, want 128000", got) } } @@ -38,8 +68,8 @@ func TestResolveModelContentLength_CompositeReference(t *testing.T) { // instance-bearing form used by tenant model instances — from the provider // catalog. func TestResolveModelContentLength_CompositeReferenceThreePart(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "", "gpt-4o@default@openai", "", ""); got != 128000 { - t.Fatalf("ResolveModelContentLength(gpt-4o@default@openai) = %d, want 128000", got) + if got := resolveTestContextLength(t.Context(), nil, "", "gpt-4o@default@openai", "", ""); got != 128000 { + t.Fatalf("resolveTestContextLength(gpt-4o@default@openai) = %d, want 128000", got) } } @@ -50,32 +80,32 @@ func TestResolveModelContentLength_CompositeReferenceThreePart(t *testing.T) { // and provider anchored on the right. When the catalog lookup fails it falls // through to the driver+modelName fallback, and to 0 when no fallback exists. func TestResolveModelContentLength_MultiSegmentComposite(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "", "a@b@c@d", "openai", "gpt-4o"); got != 128000 { - t.Fatalf("ResolveModelContentLength(a@b@c@d, openai, gpt-4o) = %d, want 128000 (driver fallback)", got) + if got := resolveTestContextLength(t.Context(), nil, "", "a@b@c@d", "openai", "gpt-4o"); got != 128000 { + t.Fatalf("resolveTestContextLength(a@b@c@d, openai, gpt-4o) = %d, want 128000 (driver fallback)", got) } - if got := ResolveModelContentLength(t.Context(), nil, "", "gpt-4o@default@openai@extra", "", ""); got != 0 { - t.Fatalf("ResolveModelContentLength(gpt-4o@default@openai@extra) = %d, want 0", got) + if got := resolveTestContextLength(t.Context(), nil, "", "gpt-4o@default@openai@extra", "", ""); got != 0 { + t.Fatalf("resolveTestContextLength(gpt-4o@default@openai@extra) = %d, want 0", got) } - modelName, instanceName, providerName, ok := splitCompositeModelRef("a@b@c@d") - if !ok || modelName != "a@b" || instanceName != "c" || providerName != "d" { - t.Fatalf("splitCompositeModelRef(a@b@c@d) = %q/%q/%q/%v, want a@b/c/d/true", modelName, instanceName, providerName, ok) + modelName, instanceName, providerName, err := common.ParseCompositeModelName("a@b@c@d") + if err != nil || modelName != "a@b" || instanceName != "c" || providerName != "d" { + t.Fatalf("ParseCompositeModelName(a@b@c@d) = %q/%q/%q/%v, want a@b/c/d", modelName, instanceName, providerName, err) } } // TestResolveModelContentLength_DriverModelFallback resolves context_length // from the resolved driver + bare model name, which needs no database. func TestResolveModelContentLength_DriverModelFallback(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "", "", "openai", "gpt-4o"); got != 128000 { - t.Fatalf("ResolveModelContentLength(openai/gpt-4o) = %d, want 128000", got) + if got := resolveTestContextLength(t.Context(), nil, "", "", "openai", "gpt-4o"); got != 128000 { + t.Fatalf("resolveTestContextLength(openai/gpt-4o) = %d, want 128000", got) } } // TestResolveModelContentLength_Unknown returns 0 for unknown references. func TestResolveModelContentLength_Unknown(t *testing.T) { - if got := ResolveModelContentLength(t.Context(), nil, "", "no-such-model@no-such-provider", "", ""); got != 0 { + if got := resolveTestContextLength(t.Context(), nil, "", "no-such-model@no-such-provider", "", ""); got != 0 { t.Fatalf("unknown model = %d, want 0", got) } - if got := ResolveModelContentLength(t.Context(), nil, "", "", "", ""); got != 0 { + if got := resolveTestContextLength(t.Context(), nil, "", "", "", ""); got != 0 { t.Fatalf("empty reference = %d, want 0", got) } } @@ -84,13 +114,12 @@ func TestResolveModelContentLength_Unknown(t *testing.T) { // tenant_model UUID row through the provider catalog. func TestResolveModelContentLength_TenantModelUUID(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, "") - if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { - t.Fatalf("ResolveModelContentLength(uuid) = %d, want 128000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { + t.Fatalf("resolveTestContextLength(uuid) = %d, want 128000", got) } } @@ -99,13 +128,12 @@ func TestResolveModelContentLength_TenantModelUUID(t *testing.T) { // provider catalog's context_length for a UUID reference. func TestResolveModelContentLength_ExtraOverrideUUID(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) - if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 32000 { - t.Fatalf("ResolveModelContentLength(uuid+extra) = %d, want 32000 (custom override wins over catalog 128000)", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 32000 { + t.Fatalf("resolveTestContextLength(uuid+extra) = %d, want 32000 (custom override wins over catalog 128000)", got) } } @@ -114,13 +142,12 @@ func TestResolveModelContentLength_ExtraOverrideUUID(t *testing.T) { // tenant's provider/instance/model rows. func TestResolveModelContentLength_ExtraOverrideComposite(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) - if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 32000 { - t.Fatalf("ResolveModelContentLength(composite+extra) = %d, want 32000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 32000 { + t.Fatalf("resolveTestContextLength(composite+extra) = %d, want 32000", got) } } @@ -130,13 +157,12 @@ func TestResolveModelContentLength_ExtraOverrideComposite(t *testing.T) { // override (the catalog cannot provide a value). func TestResolveModelContentLength_CustomModelExtraComposite(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedCustomChatModel(t, db, "my-local-model", `{"max_tokens": 32000}`) - if got := ResolveModelContentLength(ctx, db, "tenant-1", "my-local-model@OpenAI", "", ""); got != 32000 { - t.Fatalf("ResolveModelContentLength(custom+extra) = %d, want 32000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "my-local-model@OpenAI", "", ""); got != 32000 { + t.Fatalf("resolveTestContextLength(custom+extra) = %d, want 32000", got) } // Without the override the custom model is unknown to the catalog → 0. if err := db.Model(&entity.TenantModel{}). @@ -144,7 +170,7 @@ func TestResolveModelContentLength_CustomModelExtraComposite(t *testing.T) { Update("extra", "").Error; err != nil { t.Fatalf("clear custom model extra: %v", err) } - if got := ResolveModelContentLength(ctx, db, "tenant-1", "my-local-model@OpenAI", "", ""); got != 0 { + if got := resolveTestContextLength(ctx, db, "tenant-1", "my-local-model@OpenAI", "", ""); got != 0 { t.Fatalf("custom model without extra = %d, want 0", got) } } @@ -155,7 +181,6 @@ func TestResolveModelContentLength_CustomModelExtraComposite(t *testing.T) { // catalog by reference parts (gpt-4o → 128000). func TestResolveModelContentLength_InactiveInstanceCompositeFallsBackToCatalog(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) @@ -165,7 +190,7 @@ func TestResolveModelContentLength_InactiveInstanceCompositeFallsBackToCatalog(t t.Fatalf("set instance inactive: %v", err) } - if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { + if got := resolveTestContextLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { t.Fatalf("inactive-instance composite = %d, want catalog 128000", got) } } @@ -176,7 +201,6 @@ func TestResolveModelContentLength_InactiveInstanceCompositeFallsBackToCatalog(t // catalog by reference parts (gpt-4o → 128000). func TestResolveModelContentLength_InactiveModelCompositeFallsBackToCatalog(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) @@ -186,7 +210,7 @@ func TestResolveModelContentLength_InactiveModelCompositeFallsBackToCatalog(t *t t.Fatalf("set model inactive: %v", err) } - if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { + if got := resolveTestContextLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { t.Fatalf("inactive-model composite = %d, want catalog 128000", got) } } @@ -196,13 +220,12 @@ func TestResolveModelContentLength_InactiveModelCompositeFallsBackToCatalog(t *t // callers persist extra.max_tokens as a string). func TestResolveModelContentLength_ExtraOverrideStringForm(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{"max_tokens": "32000"}`) - if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 32000 { - t.Fatalf("ResolveModelContentLength(string-form extra) = %d, want 32000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 32000 { + t.Fatalf("resolveTestContextLength(string-form extra) = %d, want 32000", got) } } @@ -210,13 +233,12 @@ func TestResolveModelContentLength_ExtraOverrideStringForm(t *testing.T) { // verifies that a non-numeric max_tokens string is treated as "no override". func TestResolveModelContentLength_NonNumericStringExtraFallsBackToCatalog(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{"max_tokens": "huge"}`) - if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { - t.Fatalf("ResolveModelContentLength(non-numeric string extra) = %d, want catalog 128000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { + t.Fatalf("resolveTestContextLength(non-numeric string extra) = %d, want catalog 128000", got) } } @@ -225,13 +247,12 @@ func TestResolveModelContentLength_NonNumericStringExtraFallsBackToCatalog(t *te // context_length is used. func TestResolveModelContentLength_InvalidExtraFallsBackToCatalog(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{bad json`) - if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { - t.Fatalf("ResolveModelContentLength(invalid extra) = %d, want catalog 128000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { + t.Fatalf("resolveTestContextLength(invalid extra) = %d, want catalog 128000", got) } } @@ -239,13 +260,12 @@ func TestResolveModelContentLength_InvalidExtraFallsBackToCatalog(t *testing.T) // non-positive max_tokens override is ignored in favor of the catalog. func TestResolveModelContentLength_ZeroExtraFallsBackToCatalog(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{"max_tokens": 0}`) - if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { - t.Fatalf("ResolveModelContentLength(zero extra) = %d, want catalog 128000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 128000 { + t.Fatalf("resolveTestContextLength(zero extra) = %d, want catalog 128000", got) } } @@ -254,7 +274,6 @@ func TestResolveModelContentLength_ZeroExtraFallsBackToCatalog(t *testing.T) { // it carries a max_tokens override. func TestResolveModelContentLength_InactiveTenantModel(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() seedOpenAIChatModel(t, db, `{"max_tokens": 32000}`) @@ -266,7 +285,7 @@ func TestResolveModelContentLength_InactiveTenantModel(t *testing.T) { } // The UUID is not a composite ref, so an inactive row yields 0. - if got := ResolveModelContentLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 0 { + if got := resolveTestContextLength(ctx, db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != 0 { t.Fatalf("inactive tenant model = %d, want 0", got) } } @@ -288,8 +307,10 @@ func seedCustomChatModel(t *testing.T, db *gorm.DB, modelName, extra string) { func seedChatModel(t *testing.T, db *gorm.DB, providerID, providerName, modelName, extra string) { t.Helper() + seedTenantOwner(t, db) if err := db.Create(&entity.TenantModelProvider{ ID: providerID, + UserID: "tenant-1", ProviderName: providerName, TenantID: "tenant-1", }).Error; err != nil { @@ -316,13 +337,27 @@ func seedChatModel(t *testing.T, db *gorm.DB, providerID, providerName, modelNam } } +func seedTenantOwner(t *testing.T, db *gorm.DB) { + t.Helper() + status := entity.TenantMembershipStatusJoined + if err := db.Create(&entity.UserTenant{ + ID: "membership-tenant-1", + UserID: "tenant-1", + TenantID: "tenant-1", + Role: "owner", + Status: &status, + }).Error; err != nil { + t.Fatalf("create tenant owner: %v", err) + } +} + func openModelContextTestDB(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true}) if err != nil { t.Fatalf("open sqlite: %v", err) } - if err := db.AutoMigrate(&entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { + if err := db.AutoMigrate(&entity.UserTenant{}, &entity.TenantModelProvider{}, &entity.TenantModelInstance{}, &entity.TenantModel{}); err != nil { t.Fatalf("migrate: %v", err) } return db @@ -335,17 +370,16 @@ func openModelContextTestDB(t *testing.T) *gorm.DB { func TestResolveModelContentLength_StoredMistralOCRProvider(t *testing.T) { const want = 262144 // mistral-large-latest context_length in conf/models/mistral.json - if got := ResolveModelContentLength(t.Context(), nil, "", "mistral-large-latest@default@Mistral OCR", "", ""); got != want { + if got := resolveTestContextLength(t.Context(), nil, "", "mistral-large-latest@default@Mistral OCR", "", ""); got != want { t.Fatalf("composite without tenant row = %d, want %d", got, want) } - if got := ResolveModelContentLength(t.Context(), nil, "", "", "Mistral OCR", "mistral-large-latest"); got != want { + if got := resolveTestContextLength(t.Context(), nil, "", "", "Mistral OCR", "mistral-large-latest"); got != want { t.Fatalf("driver fallback = %d, want %d", got, want) } db := openModelContextTestDB(t) - pushDB(t, db) seedChatModel(t, db, "provider-mistral-ocr", "Mistral OCR", "mistral-large-latest", "") - if got := ResolveModelContentLength(t.Context(), db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != want { + if got := resolveTestContextLength(t.Context(), db, "tenant-1", "0123456789abcdef0123456789abcdef", "", ""); got != want { t.Fatalf("tenant model id = %d, want %d", got, want) } } @@ -359,16 +393,17 @@ func TestResolveModelContentLength_StoredMistralOCRProvider(t *testing.T) { // "model@instance@provider" reference keeps resolving that instance. func TestResolveModelContentLength_LegacyCompositeNeverUsesNonDefaultInstance(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() if err := db.Create(&entity.TenantModelProvider{ ID: "provider-openai", + UserID: "tenant-1", ProviderName: "OpenAI", TenantID: "tenant-1", }).Error; err != nil { t.Fatalf("create provider: %v", err) } + seedTenantOwner(t, db) if err := db.Create(&entity.TenantModelInstance{ ID: "instance-prod", ProviderID: "provider-openai", @@ -391,12 +426,12 @@ func TestResolveModelContentLength_LegacyCompositeNeverUsesNonDefaultInstance(t // The implicit "default" instance does not exist, so the "prod" instance's // tenant override must not answer a two-part reference. - if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { - t.Fatalf("ResolveModelContentLength(gpt-4o@OpenAI) = %d, want 128000 (the default instance is unavailable; another instance must not be substituted)", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { + t.Fatalf("resolveTestContextLength(gpt-4o@OpenAI) = %d, want 128000 (the default instance is unavailable; another instance must not be substituted)", got) } // An explicit instance reference still resolves that instance. - if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@prod@OpenAI", "", ""); got != 32000 { - t.Fatalf("ResolveModelContentLength(gpt-4o@prod@OpenAI) = %d, want 32000 (explicit instance reference)", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "gpt-4o@prod@OpenAI", "", ""); got != 32000 { + t.Fatalf("resolveTestContextLength(gpt-4o@prod@OpenAI) = %d, want 32000 (explicit instance reference)", got) } } @@ -406,16 +441,17 @@ func TestResolveModelContentLength_LegacyCompositeNeverUsesNonDefaultInstance(t // instance keeps its own override for explicit references. func TestResolveModelContentLength_LegacyCompositeDoesNotGuessBetweenActiveInstances(t *testing.T) { db := openModelContextTestDB(t) - pushDB(t, db) ctx := t.Context() if err := db.Create(&entity.TenantModelProvider{ ID: "provider-openai", + UserID: "tenant-1", ProviderName: "OpenAI", TenantID: "tenant-1", }).Error; err != nil { t.Fatalf("create provider: %v", err) } + seedTenantOwner(t, db) for _, instance := range []*entity.TenantModelInstance{ {ID: "instance-prod", ProviderID: "provider-openai", InstanceName: "prod", Status: "active"}, {ID: "instance-staging", ProviderID: "provider-openai", InstanceName: "staging", Status: "active"}, @@ -441,13 +477,13 @@ func TestResolveModelContentLength_LegacyCompositeDoesNotGuessBetweenActiveInsta } } - if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { - t.Fatalf("ResolveModelContentLength(gpt-4o@OpenAI) = %d, want 128000 (no default instance; no instance may be guessed)", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "gpt-4o@OpenAI", "", ""); got != 128000 { + t.Fatalf("resolveTestContextLength(gpt-4o@OpenAI) = %d, want 128000 (no default instance; no instance may be guessed)", got) } - if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@prod@OpenAI", "", ""); got != 32000 { - t.Fatalf("ResolveModelContentLength(gpt-4o@prod@OpenAI) = %d, want 32000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "gpt-4o@prod@OpenAI", "", ""); got != 32000 { + t.Fatalf("resolveTestContextLength(gpt-4o@prod@OpenAI) = %d, want 32000", got) } - if got := ResolveModelContentLength(ctx, db, "tenant-1", "gpt-4o@staging@OpenAI", "", ""); got != 16000 { - t.Fatalf("ResolveModelContentLength(gpt-4o@staging@OpenAI) = %d, want 16000", got) + if got := resolveTestContextLength(ctx, db, "tenant-1", "gpt-4o@staging@OpenAI", "", ""); got != 16000 { + t.Fatalf("resolveTestContextLength(gpt-4o@staging@OpenAI) = %d, want 16000", got) } } diff --git a/internal/permission/policy_ce.go b/internal/permission/policy.go similarity index 85% rename from internal/permission/policy_ce.go rename to internal/permission/policy.go index 95207324d3..c04674e847 100644 --- a/internal/permission/policy_ce.go +++ b/internal/permission/policy.go @@ -172,6 +172,9 @@ func (p *communityPolicy) Scope(ctx context.Context, subject Subject, query Scop } func (p *communityPolicy) QueryScope(ctx context.Context, subject Subject, query ScopeQuery) (QueryScope, bool, error) { + if query.Kind == ResourceKindModelProvider || query.Kind == ResourceKindModelInstance || query.Kind == ResourceKindModel { + return p.modelQueryScope(ctx, subject, query) + } if query.Kind == ResourceKindMCPServer && query.Operation == OperationRead && query.Parent.ID == "" && query.Parent.Kind == "" && query.Entry == nil { return p.mcpServerReadQueryScope(ctx, subject) } @@ -218,6 +221,93 @@ func (p *communityPolicy) QueryScope(ctx context.Context, subject Subject, query }}, true, nil } +func (p *communityPolicy) modelQueryScope(ctx context.Context, subject Subject, query ScopeQuery) (QueryScope, bool, error) { + parent := query.Parent + if query.Entry != nil || (parent.ID == "") != (parent.Kind == "") { + return QueryScope{}, false, nil + } + if err := validateSubject(subject); err != nil { + return QueryScope{}, true, err + } + source, ok := p.source.(*databaseSource) + if !ok || source.db == nil || tenantRoleMembershipEnabled(source.db) { + return QueryScope{}, false, nil + } + if query.Kind == ResourceKindModelInstance && parent.ID != "" && parent.Kind != ResourceKindModelProvider { + return QueryScope{}, false, nil + } + if query.Kind == ResourceKindModel && parent.ID != "" && parent.Kind != ResourceKindModelProvider && parent.Kind != ResourceKindModelInstance { + return QueryScope{}, false, nil + } + + providerAccess := func(providerIDColumn, parentID string, ownerOnly bool) *gorm.DB { + provider := source.db.WithContext(ctx).Model(&entity.TenantModelProvider{}). + Select("1"). + Where("tenant_model_provider.id = " + providerIDColumn) + if parentID != "" { + provider = provider.Where("tenant_model_provider.id = ?", parentID) + } + owner := source.db.WithContext(ctx).Model(&entity.UserTenant{}). + Select("user_tenant.user_id"). + Where("user_tenant.tenant_id = tenant_model_provider.tenant_id AND user_tenant.role = ? AND user_tenant.status = ?", string(RoleOwner), string(entity.StatusValid)). + Limit(1) + provider = provider.Where("tenant_model_provider.user_id = COALESCE((?), tenant_model_provider.tenant_id)", owner) + if ownerOnly { + return provider.Where("tenant_model_provider.user_id = ?", subject.UserID) + } + membership := source.db.WithContext(ctx).Model(&entity.UserTenant{}). + Select("1"). + Where("user_tenant.tenant_id = tenant_model_provider.tenant_id AND user_tenant.user_id = ? AND user_tenant.status = ? AND user_tenant.role IN ?", + subject.UserID, string(entity.StatusValid), []string{string(RoleOwner), string(RoleAdmin), string(RoleNormal)}) + return provider.Where("(tenant_model_provider.user_id = ? OR EXISTS (?))", subject.UserID, membership) + } + + switch query.Kind { + case ResourceKindModelProvider: + if parent.ID != "" || query.Operation != OperationRead { + return QueryScope{}, false, nil + } + return QueryScope{apply: func(query *gorm.DB) *gorm.DB { + db := source.db.WithContext(ctx) + owner := db.Model(&entity.UserTenant{}). + Select("user_tenant.user_id"). + Where("user_tenant.tenant_id = tenant_model_provider.tenant_id AND user_tenant.role = ? AND user_tenant.status = ?", string(RoleOwner), string(entity.StatusValid)). + Limit(1) + membership := db.Model(&entity.UserTenant{}). + Select("1"). + Where("user_tenant.tenant_id = tenant_model_provider.tenant_id AND user_tenant.user_id = ? AND user_tenant.status = ? AND user_tenant.role IN ?", + subject.UserID, string(entity.StatusValid), []string{string(RoleOwner), string(RoleAdmin), string(RoleNormal)}) + return query. + Where("tenant_model_provider.user_id = COALESCE((?), tenant_model_provider.tenant_id)", owner). + Where("(tenant_model_provider.user_id = ? OR EXISTS (?))", subject.UserID, membership) + }}, true, nil + case ResourceKindModelInstance: + if query.Operation != OperationRead && query.Operation != OperationUse { + return QueryScope{}, false, nil + } + ownerOnly := query.Operation == OperationRead + return QueryScope{apply: func(query *gorm.DB) *gorm.DB { + return query.Where("EXISTS (?)", providerAccess("tenant_model_instance.provider_id", parent.ID, ownerOnly)) + }}, true, nil + case ResourceKindModel: + if query.Operation != OperationRead && query.Operation != OperationUse { + return QueryScope{}, false, nil + } + return QueryScope{apply: func(query *gorm.DB) *gorm.DB { + modelQuery := query.Where("EXISTS (?)", providerAccess("tenant_model.provider_id", "", false)) + switch parent.Kind { + case ResourceKindModelProvider: + modelQuery = modelQuery.Where("tenant_model.provider_id = ?", parent.ID) + case ResourceKindModelInstance: + modelQuery = modelQuery.Where("tenant_model.instance_id = ?", parent.ID) + } + return modelQuery + }}, true, nil + default: + return QueryScope{}, false, nil + } +} + func (p *communityPolicy) mcpServerReadQueryScope(ctx context.Context, subject Subject) (QueryScope, bool, error) { if err := validateSubject(subject); err != nil { return QueryScope{}, true, err diff --git a/internal/permission/policy_ce_test.go b/internal/permission/policy_test.go similarity index 100% rename from internal/permission/policy_ce_test.go rename to internal/permission/policy_test.go diff --git a/internal/router/router.go b/internal/router/router.go index d458d3584d..8ebf5ec38c 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -554,9 +554,9 @@ func (r *Router) Setup(engine *gin.Engine) { model := v1.Group("/models", ToggleModel()) { - // GET /models returns the tenant's added models across - // all instances. Front-end useFetchAllAddedModels consumes this. - model.GET("", r.providerHandler.ListTenantAddedModels) + // GET /models returns models available to the authenticated user + // in the selected tenant context. + model.GET("", r.providerHandler.ListModelsForUserInTenant) model.PATCH("", r.tenantHandler.SetModels) // Tenant default-model selection (used by the agent page's useFetchDefaultModels hook) model.GET("/default", r.tenantHandler.GetDefaultModels) diff --git a/internal/server/api/server.go b/internal/server/api/server.go index a851b67e41..7c5fa85308 100644 --- a/internal/server/api/server.go +++ b/internal/server/api/server.go @@ -106,6 +106,7 @@ func startServer(ctx context.Context, serverName string, arguments *bootstrap.Se // Initialize service layer userService := service.NewUserService() + userTenantService := service.NewUserTenantService() documentService := document.NewDocumentService() datasetsService := dataset.NewDatasetService() metadataService := service.NewMetadataService() @@ -120,13 +121,12 @@ func startServer(ctx context.Context, serverName string, arguments *bootstrap.Se statsService := service.NewStatsService() connectorService := service.NewConnectorService() searchService := service.NewSearchService() - searchService.SetTenantService(tenantService) fileService := file.NewFileService(documentService) memoryService := service.NewMemoryService(dao.DB, engine.Get(), service.WithMemoryKvrocks(kvrocks.Get()), service.WithMemoryMQEngine(engine.GetMessageQueueEngine())) mcpService := service.NewMCPService() modelProviderService := service.NewModelProviderService() modelFactory := service.NewModelFactory() - modelCallService := service.NewModelCallServiceWithProviderService(modelProviderService, modelFactory) + modelCallService := service.NewModelCallService() // Wire the real MemorySaver so the Message component can persist // conversation turns to memory stores declared in the canvas DSL. @@ -154,21 +154,21 @@ func startServer(ctx context.Context, serverName string, arguments *bootstrap.Se documentDAO, agenttool.RetrievalModelConstructors{ Embedding: func(ctx context.Context, tenantID, modelRef string) (*modelModule.EmbeddingModel, error) { - access := service.ModelAccess{TenantID: tenantID} + access := service.Identity{TenantID: tenantID} if strings.TrimSpace(modelRef) == "" { return modelFactory.NewDefaultEmbeddingModel(ctx, access) } return modelFactory.NewEmbeddingModel(ctx, access, modelRef) }, Chat: func(ctx context.Context, tenantID, modelRef string) (*modelModule.ChatModel, error) { - access := service.ModelAccess{TenantID: tenantID} + access := service.Identity{TenantID: tenantID} if strings.TrimSpace(modelRef) == "" { return modelFactory.NewDefaultChatModel(ctx, access) } return modelFactory.NewChatModel(ctx, access, modelRef) }, Rerank: func(ctx context.Context, tenantID, modelRef string) (*modelModule.RerankModel, error) { - return modelFactory.NewRerankModel(ctx, service.ModelAccess{TenantID: tenantID}, modelRef) + return modelFactory.NewRerankModel(ctx, service.Identity{TenantID: tenantID}, modelRef) }, }, retrievalEnhancer, @@ -202,8 +202,8 @@ func startServer(ctx context.Context, serverName string, arguments *bootstrap.Se // Initialize handler layer authHandler := handler.NewAuthHandler() - userHandler := handler.NewUserHandler(userService) - tenantHandler := handler.NewTenantHandler(tenantService, userService, datasetsService) + userHandler := handler.NewUserHandler(userService, userTenantService) + tenantHandler := handler.NewTenantHandler(tenantService, userTenantService, datasetsService) documentHandler := handler.NewDocumentHandler(documentService, datasetsService, fileService) datasetsHandler := handler.NewDatasetsHandler(datasetsService, metadataService) systemHandler := handler.NewSystemHandler(systemService) diff --git a/internal/service/agent.go b/internal/service/agent.go index 0b17e5af18..2e690a7da8 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -2046,7 +2046,7 @@ func validateAgentChatModelReferences(ctx context.Context, userID, tenantID stri if !ok { modelRef, _ = node.Obj.Params["llm_id"].(string) } - if _, err := modelFactory.ResolveInfo(ctx, ModelAccess{UserID: userID, TenantID: tenantID}, entity.ModelTypeChat, modelRef); err != nil { + if _, err := modelFactory.ResolveInfo(ctx, Identity{UserID: userID, TenantID: tenantID}, entity.ModelTypeChat, modelRef); err != nil { if draft && (errors.Is(err, errModelConfigUnavailable) || errors.Is(err, gorm.ErrRecordNotFound)) && !errors.Is(err, permission.ErrPermissionDenied) && !errors.Is(err, permission.ErrResourceNotFound) { continue } diff --git a/internal/service/agentic_pipeline.go b/internal/service/agentic_pipeline.go index 60abec0511..1994b44e79 100644 --- a/internal/service/agentic_pipeline.go +++ b/internal/service/agentic_pipeline.go @@ -415,7 +415,7 @@ func agenticUsageFields(ctx context.Context, usage *tokenizer.RunUsage) []zap.Fi // model choice to every chat model owned by the tenant: a failover chain is // exactly the list the dialog's author chose, and nothing else. func (s *ChatPipelineService) agenticModelChain(ctx context.Context, userID string, chat *entity.Chat) ([]*modelModule.ChatModel, error) { - access := ModelAccess{UserID: userID, TenantID: chat.TenantID} + access := Identity{UserID: userID, TenantID: chat.TenantID} primary, err := s.agenticPrimaryModel(ctx, access, chat) if err != nil { return nil, err @@ -445,7 +445,7 @@ func (s *ChatPipelineService) agenticModelChain(ctx context.Context, userID stri // agenticPrimaryModel resolves the dialog's own model, or the tenant default // when it selected none. -func (s *ChatPipelineService) agenticPrimaryModel(ctx context.Context, access ModelAccess, chat *entity.Chat) (*modelModule.ChatModel, error) { +func (s *ChatPipelineService) agenticPrimaryModel(ctx context.Context, access Identity, chat *entity.Chat) (*modelModule.ChatModel, error) { var ( model *modelModule.ChatModel err error diff --git a/internal/service/ask_service.go b/internal/service/ask_service.go index b0f66f4372..53fac41736 100644 --- a/internal/service/ask_service.go +++ b/internal/service/ask_service.go @@ -81,7 +81,7 @@ type AskStreamOptions struct { // Retriever abstracts chunk retrieval for AskService. type Retriever interface { - RetrievalTest(ctx context.Context, req *RetrievalTestRequest, userID string) (*RetrievalTestResponse, error) + RetrievalTest(ctx context.Context, req *RetrievalTestRequest) (*RetrievalTestResponse, error) } // StreamingLLM abstracts streaming chat for AskService. The first channel @@ -133,22 +133,22 @@ func NewAskService(retriever Retriever, embedder Embedder, tokenBudget, minStrea // Stream runs the full ask pipeline. llm must not be nil. The returned // channel is closed when the pipeline completes or ctx is cancelled. -func (s *AskService) Stream(ctx context.Context, llm StreamingLLM, userID, question string, kbIDs []string) <-chan AskDelta { - return s.StreamWithOptions(ctx, llm, userID, question, kbIDs, AskStreamOptions{}) +func (s *AskService) Stream(ctx context.Context, llm StreamingLLM, question string, kbIDs []string) <-chan AskDelta { + return s.StreamWithOptions(ctx, llm, question, kbIDs, AskStreamOptions{}) } // StreamWithOptions runs Stream while allowing callers such as saved Search // apps to pass search_config retrieval options through to RetrievalTest. -func (s *AskService) StreamWithOptions(ctx context.Context, llm StreamingLLM, userID, question string, datasetIDs []string, opts AskStreamOptions) <-chan AskDelta { +func (s *AskService) StreamWithOptions(ctx context.Context, llm StreamingLLM, question string, datasetIDs []string, opts AskStreamOptions) <-chan AskDelta { out := make(chan AskDelta, 32) go func() { defer close(out) - s.run(ctx, llm, userID, question, datasetIDs, opts, out) + s.run(ctx, llm, question, datasetIDs, opts, out) }() return out } -func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question string, datasetIDs []string, opts AskStreamOptions, out chan<- AskDelta) { +func (s *AskService) run(ctx context.Context, llm StreamingLLM, question string, datasetIDs []string, opts AskStreamOptions, out chan<- AskDelta) { // Phase 1: Retrieval. topK := DefaultAskTopK if opts.TopK != nil { @@ -190,7 +190,7 @@ func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question req.Page = &page req.Size = &ps - result, err := s.retriever.RetrievalTest(ctx, req, userID) + result, err := s.retriever.RetrievalTest(ctx, req) if err != nil { common.Warn("AskService retrieval failed", zap.Error(err)) s.sendOrCancel(out, AskDelta{Kind: AskDeltaError, Value: "retrieval failed"}, ctx) diff --git a/internal/service/ask_service_test.go b/internal/service/ask_service_test.go index ee36610644..8e445e43dc 100644 --- a/internal/service/ask_service_test.go +++ b/internal/service/ask_service_test.go @@ -32,7 +32,7 @@ type fakeRetriever struct { err error } -func (r *fakeRetriever) RetrievalTest(ctx context.Context, req *RetrievalTestRequest, userID string) (*RetrievalTestResponse, error) { +func (r *fakeRetriever) RetrievalTest(ctx context.Context, req *RetrievalTestRequest) (*RetrievalTestResponse, error) { if r.err != nil { return nil, r.err } @@ -98,7 +98,7 @@ func TestAskService_RetrievalError(t *testing.T) { ret := &fakeRetriever{err: fmt.Errorf("engine down")} llm := &fakeStreamLLM{chunks: []string{"answer"}} svc := NewAskService(ret, nil, 0, 0) - deltas := collect(svc.Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(t.Context(), llm, "test", []string{"kb1"})) if len(deltas) < 1 || deltas[0].Kind != AskDeltaError { t.Fatalf("expected error delta, got %+v", deltas) } @@ -108,7 +108,7 @@ func TestAskService_EmptyResult(t *testing.T) { ret := &fakeRetriever{result: &RetrievalTestResponse{Chunks: []map[string]interface{}{}}} llm := &fakeStreamLLM{chunks: []string{"answer"}} svc := NewAskService(ret, nil, 0, 0) - deltas := collect(svc.Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(t.Context(), llm, "test", []string{"kb1"})) if len(deltas) < 1 || !strings.Contains(deltas[0].Value, "no relevant information") { t.Fatalf("expected 'no relevant information', got %+v", deltas) } @@ -123,7 +123,7 @@ func TestAskService_StreamingFlow(t *testing.T) { }} llm := &fakeStreamLLM{chunks: []string{"Hello", " world"}} svc := NewAskService(ret, nil, 0, 0) - deltas := collect(svc.Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(t.Context(), llm, "test", []string{"kb1"})) var hasAnswer, hasFinal bool for _, d := range deltas { @@ -150,7 +150,7 @@ func TestAskService_EmptyLLMStreamReturnsError(t *testing.T) { }} llm := &fakeStreamLLM{} svc := NewAskService(ret, nil, 0, 0) - deltas := collect(svc.Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(t.Context(), llm, "test", []string{"kb1"})) if len(deltas) != 1 || deltas[0].Kind != AskDeltaError || !strings.Contains(deltas[0].Value, "LLM call failed") { t.Fatalf("expected LLM error delta, got %+v", deltas) @@ -166,7 +166,7 @@ func TestAskService_ThinkTags(t *testing.T) { }} llm := &fakeStreamLLM{chunks: []string{"", "reasoning...", "", "visible answer"}} svc := NewAskService(ret, nil, 0, 0) - deltas := collect(svc.Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(t.Context(), llm, "test", []string{"kb1"})) var hasMarker bool for _, d := range deltas { @@ -187,7 +187,7 @@ func TestAskService_LLMError(t *testing.T) { }} llm := &fakeStreamLLM{err: fmt.Errorf("model offline")} svc := NewAskService(ret, nil, 0, 0) - deltas := collect(svc.Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(t.Context(), llm, "test", []string{"kb1"})) if len(deltas) < 1 || deltas[0].Kind != AskDeltaError { t.Fatalf("expected error delta, got %+v", deltas) } @@ -209,7 +209,7 @@ func TestAskService_TemperatureFromLLMSetting(t *testing.T) { "temperature_enabled": true, }, }) - deltas := collect(svc.StreamWithOptions(t.Context(), llm, "user1", "test", []string{"kb1"}, opts)) + deltas := collect(svc.StreamWithOptions(t.Context(), llm, "test", []string{"kb1"}, opts)) var hasFinal bool for _, d := range deltas { @@ -230,7 +230,7 @@ func TestAskService_DisablesThinkingForSummary(t *testing.T) { Chunks: []map[string]interface{}{{"id": "c1", "content_with_weight": "chunk", "docnm_kwd": "Doc", "kb_id": "kb1", "doc_id": "d1"}}, }} llm := &capturingStreamLLM{fakeStreamLLM: fakeStreamLLM{chunks: []string{"answer"}}} - collect(NewAskService(ret, nil, 0, 0).Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + collect(NewAskService(ret, nil, 0, 0).Stream(t.Context(), llm, "test", []string{"kb1"})) if llm.gotConfig == nil || llm.gotConfig.Thinking == nil || *llm.gotConfig.Thinking { t.Fatalf("summary must explicitly disable thinking, got %+v", llm.gotConfig) } @@ -252,7 +252,7 @@ func TestAskService_TemperatureDisabledUsesDefault(t *testing.T) { "temperature_enabled": false, }, }) - deltas := collect(svc.StreamWithOptions(t.Context(), llm, "user1", "test", []string{"kb1"}, opts)) + deltas := collect(svc.StreamWithOptions(t.Context(), llm, "test", []string{"kb1"}, opts)) var hasFinal bool for _, d := range deltas { @@ -285,7 +285,7 @@ func TestAskService_TemperatureMissingUsesLLMSettingDefaults(t *testing.T) { "temperature_enabled": true, }, }) - collect(svc.StreamWithOptions(t.Context(), llm, "user1", "test", []string{"kb1"}, opts)) + collect(svc.StreamWithOptions(t.Context(), llm, "test", []string{"kb1"}, opts)) if llm.gotConfig == nil || llm.gotConfig.Temperature == nil || *llm.gotConfig.Temperature != DefaultAskTemperature { t.Errorf("expected LLM_SETTING_DEFAULTS temperature %v, got %+v", DefaultAskTemperature, llm.gotConfig) @@ -373,7 +373,7 @@ func TestAskService_ContextCancel(t *testing.T) { svc := NewAskService(ret, nil, 0, 0) ctx, cancel := context.WithCancel(context.Background()) cancel() // cancel immediately - deltas := collect(svc.Stream(ctx, llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(ctx, llm, "test", []string{"kb1"})) // Should get no deltas (or very few) since context is cancelled. _ = deltas } @@ -387,7 +387,7 @@ func TestAskService_AsyncLLMErrorSurfaces(t *testing.T) { }} llm := &asyncErrStreamLLM{fakeStreamLLM: fakeStreamLLM{chunks: []string{"partial"}}} svc := NewAskService(ret, nil, 0, 0) - deltas := collect(svc.Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(t.Context(), llm, "test", []string{"kb1"})) var errDelta *AskDelta for i := range deltas { @@ -412,7 +412,7 @@ func TestAskService_LLMErrorValueIsRealError(t *testing.T) { }} llm := &fakeStreamLLM{err: fmt.Errorf("provider name missing in model name: gpt-4o")} svc := NewAskService(ret, nil, 0, 0) - deltas := collect(svc.Stream(t.Context(), llm, "user1", "test", []string{"kb1"})) + deltas := collect(svc.Stream(t.Context(), llm, "test", []string{"kb1"})) if len(deltas) < 1 || deltas[0].Kind != AskDeltaError { t.Fatalf("expected error delta, got %+v", deltas) } diff --git a/internal/service/chat.go b/internal/service/chat.go index 2b7da9c7ce..9cec5b42cd 100644 --- a/internal/service/chat.go +++ b/internal/service/chat.go @@ -461,7 +461,7 @@ func resolveCreateLLMID(ctx context.Context, llmID, tenantID string, llmSetting } } modelFactory := NewModelFactory() - target, err := modelFactory.ResolveInfo(ctx, ModelAccess{TenantID: tenantID}, modelType, llmID) + target, err := modelFactory.ResolveInfo(ctx, Identity{TenantID: tenantID}, modelType, llmID) if err != nil { return "", fmt.Errorf("`llm_id` %s doesn't exist", llmID) } @@ -477,7 +477,7 @@ func resolveCreateRerankID(ctx context.Context, rerankID, tenantID string) (stri return "", nil } modelFactory := NewModelFactory() - target, err := modelFactory.ResolveInfo(ctx, ModelAccess{TenantID: tenantID}, entity.ModelTypeRerank, rerankID) + target, err := modelFactory.ResolveInfo(ctx, Identity{TenantID: tenantID}, entity.ModelTypeRerank, rerankID) if err != nil { return "", fmt.Errorf("`rerank_id` %s doesn't exist", rerankID) } @@ -936,7 +936,7 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, req["tenant_decision_id"] = nil } else { modelFactory := NewModelFactory() - target, err := modelFactory.ResolveInfo(ctx, ModelAccess{TenantID: userID}, entity.ModelTypeDecision, decisionID) + target, err := modelFactory.ResolveInfo(ctx, Identity{TenantID: userID}, entity.ModelTypeDecision, decisionID) if err != nil { return nil, fmt.Errorf("`decision_id` %s doesn't exist", decisionID) } @@ -1190,7 +1190,7 @@ func (s *ChatService) resolveRESTLLMID(ctx context.Context, llmID, tenantID stri } } modelFactory := NewModelFactory() - target, err := modelFactory.ResolveInfo(ctx, ModelAccess{TenantID: tenantID}, modelType, llmID) + target, err := modelFactory.ResolveInfo(ctx, Identity{TenantID: tenantID}, modelType, llmID) if err != nil { return "", fmt.Errorf("`llm_id` %s doesn't exist", llmID) } @@ -1206,7 +1206,7 @@ func (s *ChatService) resolveRESTRerankID(ctx context.Context, rerankID, tenantI return "", nil } modelFactory := NewModelFactory() - target, err := modelFactory.ResolveInfo(ctx, ModelAccess{TenantID: tenantID}, entity.ModelTypeRerank, rerankID) + target, err := modelFactory.ResolveInfo(ctx, Identity{TenantID: tenantID}, entity.ModelTypeRerank, rerankID) if err != nil { return "", fmt.Errorf("`rerank_id` %s doesn't exist", rerankID) } diff --git a/internal/service/chat_pipeline.go b/internal/service/chat_pipeline.go index 6c2109b683..98f9a1d6be 100644 --- a/internal/service/chat_pipeline.go +++ b/internal/service/chat_pipeline.go @@ -1572,7 +1572,7 @@ func (s *ChatPipelineService) AsyncChatSolo( } // 4. Construct the selected chat model through the factory. - access := ModelAccess{UserID: userID, TenantID: chat.TenantID} + access := Identity{UserID: userID, TenantID: chat.TenantID} var chatModel *modelModule.ChatModel if strings.TrimSpace(chat.LLMID) == "" { chatModel, err = s.ModelFactory.NewDefaultChatModel(ctx, access) @@ -2071,7 +2071,7 @@ func tokenizeText(text string) string { // models reject them at the provider (e.g. Zhipu GLM error 1210: // messages.content.type only allows 'text'). func (s *ChatPipelineService) getLLMModelConfig(ctx context.Context, chat *entity.Chat) (map[string]interface{}, string, string, string, error) { - access := ModelAccess{TenantID: chat.TenantID} + access := Identity{TenantID: chat.TenantID} var info *modelModule.ModelInfo var err error if strings.TrimSpace(chat.LLMID) == "" { @@ -2153,7 +2153,7 @@ func (s *ChatPipelineService) getModels(ctx context.Context, chat *entity.Chat) } if kbs[0].EmbdID != "" { embdTenantID := kbs[0].TenantID - model, err := s.ModelFactory.NewEmbeddingModel(ctx, ModelAccess{TenantID: embdTenantID}, kbs[0].EmbdID) + model, err := s.ModelFactory.NewEmbeddingModel(ctx, Identity{TenantID: embdTenantID}, kbs[0].EmbdID) if err != nil { common.Warn("Failed to get embedding model for chat retrieval", zap.String("embdID", kbs[0].EmbdID), @@ -2166,7 +2166,7 @@ func (s *ChatPipelineService) getModels(ctx context.Context, chat *entity.Chat) } // Chat model. - access := ModelAccess{TenantID: chat.TenantID} + access := Identity{TenantID: chat.TenantID} var chatModel *modelModule.ChatModel if strings.TrimSpace(chat.LLMID) == "" { chatModel, _ = s.ModelFactory.NewDefaultChatModel(ctx, access) @@ -2790,7 +2790,7 @@ func (s *ChatPipelineService) buildChatDriver(ctx context.Context, chat *entity. if chatModel != nil { return chatModel } - access := ModelAccess{TenantID: chat.TenantID} + access := Identity{TenantID: chat.TenantID} if strings.TrimSpace(chat.LLMID) == "" { model, err := s.ModelFactory.NewDefaultChatModel(ctx, access) if err != nil { diff --git a/internal/service/chat_pipeline_test.go b/internal/service/chat_pipeline_test.go index ea6cfa605c..c38e4b4340 100644 --- a/internal/service/chat_pipeline_test.go +++ b/internal/service/chat_pipeline_test.go @@ -1862,7 +1862,7 @@ func setupChatPipelineToolSupportTestDB(t *testing.T) *gorm.DB { activeStatus := "1" rows := []interface{}{ &entity.UserTenant{ID: "user-tenant-1", UserID: "user-1", TenantID: "tenant-1", Role: "owner", InvitedBy: "user-1", Status: &activeStatus}, - &entity.TenantModelProvider{ID: "provider-zhipu", TenantID: "tenant-1", ProviderName: "ZHIPU-AI"}, + &entity.TenantModelProvider{ID: "provider-zhipu", UserID: "user-1", TenantID: "tenant-1", ProviderName: "ZHIPU-AI"}, &entity.TenantModelInstance{ID: "instance-zhipu", ProviderID: "provider-zhipu", InstanceName: "default", APIKey: "sk-test", Status: "active", Extra: "{}"}, &entity.TenantModel{ID: "model-tools-on", ProviderID: "provider-zhipu", InstanceID: "instance-zhipu", ModelName: "glm-4-plus", ModelType: int(entity.ModelTypeChat), Status: "active", Extra: `{"is_tools":true}`}, &entity.TenantModel{ID: "model-tools-off", ProviderID: "provider-zhipu", InstanceID: "instance-zhipu", ModelName: "glm-4-flash", ModelType: int(entity.ModelTypeChat), Status: "active", Extra: `{"is_tools":false}`}, @@ -1941,7 +1941,7 @@ func TestResolveModelConfigRejectsImage2TextOnlyModel(t *testing.T) { svc := NewChatPipelineService() _, err := svc.ModelFactory.resolveConfig( - t.Context(), ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-image2text-tools", + t.Context(), Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-image2text-tools", ) if err == nil { t.Fatal("ResolveModelConfig accepted an image2text-only model as chat") diff --git a/internal/service/chat_pipeline_vision_test.go b/internal/service/chat_pipeline_vision_test.go index 757f2209b5..6386a38720 100644 --- a/internal/service/chat_pipeline_vision_test.go +++ b/internal/service/chat_pipeline_vision_test.go @@ -43,7 +43,7 @@ func setupChatPipelineVisionTestDB(t *testing.T) *gorm.DB { activeStatus := "1" rows := []interface{}{ &entity.UserTenant{ID: "user-tenant-1", UserID: "user-1", TenantID: "tenant-1", Role: "owner", InvitedBy: "user-1", Status: &activeStatus}, - &entity.TenantModelProvider{ID: "provider-zhipu", TenantID: "tenant-1", ProviderName: "ZHIPU-AI"}, + &entity.TenantModelProvider{ID: "provider-zhipu", UserID: "user-1", TenantID: "tenant-1", ProviderName: "ZHIPU-AI"}, &entity.TenantModelInstance{ID: "instance-zhipu", ProviderID: "provider-zhipu", InstanceName: "default", APIKey: "sk-test", Status: "active", Extra: "{}"}, &entity.TenantModel{ID: "model-flash-chat", ProviderID: "provider-zhipu", InstanceID: "instance-zhipu", ModelName: "glm-4-flash", ModelType: int(entity.ModelTypeChat), Status: "active"}, &entity.TenantModel{ID: "model-4v", ProviderID: "provider-zhipu", InstanceID: "instance-zhipu", ModelName: "glm-4v", ModelType: int(entity.ModelTypeChat | entity.ModelTypeImage2Text), Status: "active"}, diff --git a/internal/service/chat_rest_update_test.go b/internal/service/chat_rest_update_test.go index 836c1fef52..fa7dfbf98a 100644 --- a/internal/service/chat_rest_update_test.go +++ b/internal/service/chat_rest_update_test.go @@ -56,6 +56,7 @@ func setupChatRESTUpdateServiceTestDB(t *testing.T) *gorm.DB { } if err := db.Create(&entity.TenantModelProvider{ ID: "provider-a", + UserID: "user-1", TenantID: "user-1", ProviderName: "OpenAI", }).Error; err != nil { diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go index 494e11fe71..4980983014 100644 --- a/internal/service/chat_session.go +++ b/internal/service/chat_session.go @@ -63,9 +63,9 @@ type chatPipelineRunner interface { } type chatModelInfoResolver interface { - ResolveInfo(ctx context.Context, access ModelAccess, modelType entity.ModelType, modelRef string) (*modelModule.ModelInfo, error) - ResolveDefaultInfo(ctx context.Context, access ModelAccess, modelType entity.ModelType) (*modelModule.ModelInfo, error) - ResolveModelType(ctx context.Context, access ModelAccess, modelRef string) ([]entity.ModelType, error) + ResolveInfo(ctx context.Context, access Identity, modelType entity.ModelType, modelRef string) (*modelModule.ModelInfo, error) + ResolveDefaultInfo(ctx context.Context, access Identity, modelType entity.ModelType) (*modelModule.ModelInfo, error) + ResolveModelType(ctx context.Context, access Identity, modelRef string) ([]entity.ModelType, error) } // chunkFeedbackApplier is the dispatch seam for chunk-level feedback @@ -1225,13 +1225,19 @@ func (s *ChatSessionService) ChatCompletions( dialog.LLMID = llmID dialog.LLMSetting = genConfig } else if dialog.LLMID == "" { - membership, err := NewTenantService().defaultModelMembership(ctx, userID, dialog.TenantID) - if err != nil || membership.LLMID == "" { + modelInfo, err := NewModelFactory().ResolveDefaultInfo(ctx, Identity{UserID: userID, TenantID: dialog.TenantID}, entity.ModelTypeChat) + if err != nil { return fail(errors.New("no default chat model for tenant")) } - dialog.LLMID = membership.LLMID - if membership.TenantLLMID != nil && *membership.TenantLLMID != "" { - dialog.TenantLLMID = membership.TenantLLMID + if modelInfo.InstanceName == "" { + dialog.LLMID = fmt.Sprintf("%s@%s", modelInfo.Name, modelInfo.ProviderName) + } else { + dialog.LLMID = fmt.Sprintf("%s@%s@%s", modelInfo.Name, modelInfo.InstanceName, modelInfo.ProviderName) + } + dialog.TenantLLMID = nil + if modelInfo.ID != "" { + modelID := modelInfo.ID + dialog.TenantLLMID = &modelID } if dialog.LLMSetting == nil { dialog.LLMSetting = entity.JSONMap{} @@ -1818,7 +1824,7 @@ func (s *ChatSessionService) checkTenantLLMAPIKey(ctx context.Context, tenantID, modelFactory = NewModelFactory() } var err error - access := ModelAccess{TenantID: tenantID} + access := Identity{TenantID: tenantID} if modelName == "" { _, err = modelFactory.ResolveDefaultInfo(ctx, access, entity.ModelTypeChat) } else { diff --git a/internal/service/chat_session_test.go b/internal/service/chat_session_test.go index befecef110..6cf6ba7c89 100644 --- a/internal/service/chat_session_test.go +++ b/internal/service/chat_session_test.go @@ -454,7 +454,7 @@ type fakeChatModelConfigResolver struct { err error } -func (f *fakeChatModelConfigResolver) ResolveInfo(_ context.Context, access ModelAccess, _ entity.ModelType, modelRef string) (*modelModule.ModelInfo, error) { +func (f *fakeChatModelConfigResolver) ResolveInfo(_ context.Context, access Identity, _ entity.ModelType, modelRef string) (*modelModule.ModelInfo, error) { f.tenantID = access.TenantID f.llmID = modelRef if f.err != nil { @@ -463,7 +463,7 @@ func (f *fakeChatModelConfigResolver) ResolveInfo(_ context.Context, access Mode return &modelModule.ModelInfo{ID: "resolved-model-id", Name: "resolved-model"}, nil } -func (f *fakeChatModelConfigResolver) ResolveDefaultInfo(_ context.Context, access ModelAccess, _ entity.ModelType) (*modelModule.ModelInfo, error) { +func (f *fakeChatModelConfigResolver) ResolveDefaultInfo(_ context.Context, access Identity, _ entity.ModelType) (*modelModule.ModelInfo, error) { f.tenantID = access.TenantID if f.err != nil { return nil, f.err @@ -471,7 +471,7 @@ func (f *fakeChatModelConfigResolver) ResolveDefaultInfo(_ context.Context, acce return &modelModule.ModelInfo{ID: "resolved-model-id", Name: "resolved-model"}, nil } -func (f *fakeChatModelConfigResolver) ResolveModelType(ctx context.Context, access ModelAccess, modelRef string) ([]entity.ModelType, error) { +func (f *fakeChatModelConfigResolver) ResolveModelType(ctx context.Context, access Identity, modelRef string) ([]entity.ModelType, error) { return []entity.ModelType{entity.ModelTypeChat}, nil } diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index ba9443459f..d853f228eb 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -136,8 +136,9 @@ func NewChunkService() *ChunkService { // - Builds doc_aggs by aggregating chunks per document // 7. knowledge graph retrieval (not implemented) // 8. Apply retrieval by children to group child chunks under parent chunks -func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest, userID string) (*service.RetrievalTestResponse, error) { - common.Info("RetrievalTest started", zap.String("userID", userID), zap.Any("kbID", req.Datasets), zap.String("question", req.Question)) +func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.RetrievalTestRequest) (*service.RetrievalTestResponse, error) { + subject, _ := permission.SubjectFromContext(ctx) + common.Info("RetrievalTest started", zap.String("userID", subject.UserID), zap.Any("kbID", req.Datasets), zap.String("question", req.Question)) common.Debug(fmt.Sprintf("RetrievalTest request:\n"+ " kbID=%v\n"+ @@ -171,7 +172,10 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval var tenantIDs []string var kbRecords []*entity.Knowledgebase for _, datasetID := range req.Datasets { - if err := s.checkDatasetAccess(ctx, datasetID, userID, permission.OperationUse); err != nil { + if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.Subject{}, permission.ResourceRef{ + Kind: permission.ResourceKindDataset, + ID: datasetID, + }, permission.OperationUse); err != nil { _, permissionErr := permissionresponse.NormalizeHidden(err) return nil, permissionErr } @@ -221,7 +225,7 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval if s.searchService == nil { return nil, fmt.Errorf("search service not configured") } - searchDetail, err := s.searchService.GetDetail(ctx, userID, *req.SearchID) + searchDetail, err := s.searchService.GetDetail(ctx, *req.SearchID) if err != nil { return nil, fmt.Errorf("failed to get search app %s: %w", *req.SearchID, err) } @@ -257,7 +261,7 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval if filter != nil { method, _ := filter["method"].(string) if method == "auto" || method == "semi_auto" { - access := service.ModelAccess{TenantID: tenantIDs[0]} + access := service.Identity{TenantID: tenantIDs[0]} if chatID != "" { chatModelForFilter, getErr = modelFactory.NewChatModel(ctx, access, chatID) if getErr != nil { @@ -305,20 +309,17 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval // Get chat model for cross_languages and keyword_extraction var llmModelName string if len(req.CrossLanguages) > 0 || (req.Keyword != nil && *req.Keyword) { - tenantSvc := service.NewTenantService() var err error - llmModelName, err = tenantSvc.GetDefaultModelNameForTenant(ctx, tenantIDs[0], entity.ModelTypeChat) - if err != nil || llmModelName == "" { + chatModel, err = modelFactory.NewDefaultChatModel(ctx, service.Identity{TenantID: tenantIDs[0]}) + if err != nil { common.Warn("Failed to get default chat model name for LLM transformations", zap.Error(err)) } else { - chatModel, err = modelFactory.NewChatModel(ctx, service.ModelAccess{TenantID: tenantIDs[0]}, llmModelName) - if err != nil { - common.Warn("Failed to get chat model for LLM transformations", zap.Error(err)) - } else { - common.Info("Fetched chat model (tenant default) for cross_languages/keyword_extraction", - zap.String("tenantID", tenantIDs[0]), - zap.String("modelName", llmModelName)) + if info := chatModel.Info(); info != nil { + llmModelName = service.ModelReference(info) } + common.Info("Fetched chat model (tenant default) for cross_languages/keyword_extraction", + zap.String("tenantID", tenantIDs[0]), + zap.String("modelName", llmModelName)) } } @@ -356,7 +357,7 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval labels := metadataSvc.LabelQuestion(ctx, modifiedQuestion, kbRecords) // Determine embedding model. - access := service.ModelAccess{TenantID: tenantIDs[0]} + access := service.Identity{TenantID: tenantIDs[0]} var embeddingModel *models.EmbeddingModel var embdID string if kbRecords[0].TenantEmbdID != nil && *kbRecords[0].TenantEmbdID != "" { @@ -445,7 +446,7 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval // Infinity/OceanBase chunks already carry real vectors and are left unchanged. hydrateChunkVectors(ctx, s.docEngine, filteredChunks, req.Datasets, tenantIDs) - common.Info("RetrievalTest completed", zap.String("userID", userID), zap.Any("kbID", req.Datasets), zap.String("question", req.Question), zap.Int64("chunkCount", int64(len(filteredChunks)))) + common.Info("RetrievalTest completed", zap.String("userID", subject.UserID), zap.Any("kbID", req.Datasets), zap.String("question", req.Question), zap.Int64("chunkCount", int64(len(filteredChunks)))) return &service.RetrievalTestResponse{ Chunks: filteredChunks, @@ -1741,7 +1742,7 @@ func (s *ChunkService) getEmbeddingModel(ctx context.Context, tenantID, embdID s if s.getEmbeddingModelFunc != nil { return s.getEmbeddingModelFunc(tenantID, embdID) } - embeddingModel, err := service.NewModelFactory().NewEmbeddingModel(ctx, service.ModelAccess{TenantID: tenantID}, embdID) + embeddingModel, err := service.NewModelFactory().NewEmbeddingModel(ctx, service.Identity{TenantID: tenantID}, embdID) if err != nil { return nil, err } diff --git a/internal/service/chunk/chunk_permissions_test.go b/internal/service/chunk/chunk_permissions_test.go index 637f9a7fa4..52bb36ad8c 100644 --- a/internal/service/chunk/chunk_permissions_test.go +++ b/internal/service/chunk/chunk_permissions_test.go @@ -25,7 +25,7 @@ func TestRetrievalTestChecksDatasetUsePermission(t *testing.T) { _, err := (&ChunkService{kbDAO: dao.NewKnowledgebaseDAO()}).RetrievalTest(ctx, &service.RetrievalTestRequest{ Question: "query", Datasets: []string{"kb-1"}, - }, "outsider") + }) if !permissionresponse.IsPermissionError(err) || !errors.Is(err, permission.ErrPermissionDenied) { t.Fatalf("RetrievalTest() error = %v, want permission denial", err) } diff --git a/internal/service/compilation_template_group_service.go b/internal/service/compilation_template_group_service.go index be034b9d52..0c35b40aa4 100644 --- a/internal/service/compilation_template_group_service.go +++ b/internal/service/compilation_template_group_service.go @@ -25,6 +25,7 @@ import ( "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity" + "ragflow/internal/permission" "gorm.io/gorm" ) @@ -269,7 +270,7 @@ func (s *CompilationTemplateGroupService) buildGroupItem(ctx context.Context, te Name: c.Name, Description: derefString(c.Description), Kind: c.Kind, - Config: fillConfigDefaultLLM(ctx, c.Config, c.TenantID), + Config: fillConfigDefaultLLM(ctx, c.Config, identityForTemplateTenant(ctx, tenantID)), CreateTime: timeStr(c.CreateTime), UpdateTime: timeStr(c.UpdateTime), }) @@ -392,7 +393,7 @@ func (s *CompilationTemplateGroupService) insertChildren(ctx context.Context, db } } desc := child.Description - config := fillConfigDefaultLLM(ctx, child.Config, &tenantID) + config := fillConfigDefaultLLM(ctx, child.Config, identityForTemplateTenant(ctx, tenantID)) if config == nil { config = entity.JSONMap{} } @@ -496,7 +497,7 @@ func (s *CompilationTemplateGroupService) reconcileChildren(ctx context.Context, // Pass 3: apply updates and inserts. for _, item := range resolved { - config := fillConfigDefaultLLM(ctx, item.child.Config, &tenantID) + config := fillConfigDefaultLLM(ctx, item.child.Config, identityForTemplateTenant(ctx, tenantID)) if config == nil { config = entity.JSONMap{} } @@ -549,6 +550,14 @@ func (s *CompilationTemplateGroupService) reconcileChildren(ctx context.Context, return s.assertUniqueChildNames(ctx, db, groupID, oldNames) } +func identityForTemplateTenant(ctx context.Context, tenantID string) Identity { + access := Identity{UserID: tenantID, TenantID: tenantID} + if subject, ok := permission.SubjectFromContext(ctx); ok && strings.TrimSpace(subject.UserID) != "" { + access.UserID = subject.UserID + } + return access +} + // assertUniqueChildNames rejects a group whose valid children collide on a // case-insensitive name, matching the DAO's duplicate guard which ignores // built-in templates. oldNames maps each child id to its name before the diff --git a/internal/service/compilation_template_service.go b/internal/service/compilation_template_service.go index 101fe88162..fefc86468d 100644 --- a/internal/service/compilation_template_service.go +++ b/internal/service/compilation_template_service.go @@ -38,15 +38,19 @@ import ( // fill_config_default_llm: when the config has no explicit llm_id it lazily // fills the current member's default chat model id (read-side only; the DB row is left // untouched). -func fillConfigDefaultLLM(ctx context.Context, config entity.JSONMap, tenantID *string) entity.JSONMap { - if len(config) == 0 || tenantID == nil || *tenantID == "" { +func fillConfigDefaultLLM(ctx context.Context, config entity.JSONMap, access Identity) entity.JSONMap { + if len(config) == 0 || access.TenantID == "" { return config } if _, ok := config["llm_id"]; ok { return config } - _, modelID, err := NewTenantService().GetDefaultModelRefsForTenant(ctx, *tenantID, entity.ModelTypeChat) - if err != nil || modelID == "" { + info, err := NewModelFactory().ResolveDefaultInfo(ctx, access, entity.ModelTypeChat) + if err != nil { + return config + } + modelID := ModelReference(info) + if modelID == "" { return config } out := make(entity.JSONMap, len(config)+1) @@ -114,14 +118,17 @@ func NewCompilationTemplateService() *CompilationTemplateService { // ListBuiltins returns the read-only template palette loaded from the YAML // definitions. Built-in templates are never persisted in the database. -func (s *CompilationTemplateService) ListBuiltins(ctx context.Context, tenantID string) ([]*BuiltinTemplate, error) { +func (s *CompilationTemplateService) ListBuiltins(ctx context.Context) ([]*BuiltinTemplate, error) { + access, err := ResolveIdentity(ctx) + if err != nil { + return nil, err + } templates, err := loadBuiltinTemplates() if err != nil { return nil, err } - tenantPtr := &tenantID for _, template := range templates { - template.Config = fillConfigDefaultLLM(ctx, template.Config, tenantPtr) + template.Config = fillConfigDefaultLLM(ctx, template.Config, access) } return s.sortBuiltins(templates), nil } diff --git a/internal/service/dataset/create_test.go b/internal/service/dataset/create_test.go index 2a769316fc..cd309ad22a 100644 --- a/internal/service/dataset/create_test.go +++ b/internal/service/dataset/create_test.go @@ -6,6 +6,7 @@ import ( "strings" "testing" + "gorm.io/gorm" "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity" @@ -30,13 +31,6 @@ func insertCreateDatasetTenant(t *testing.T, tenantID string) { } } ensureDatasetTestMembership(t, tenantID, tenantID, "owner") - if err := dao.DB.Model(&entity.UserTenant{}). - Where("user_id = ? AND tenant_id = ?", tenantID, tenantID). - Updates(map[string]interface{}{ - "llm_id": "llm-default", "embd_id": "embd-default", "tenant_embd_id": "embd-1", "asr_id": "asr-default", - }).Error; err != nil { - t.Fatalf("configure owner model defaults: %v", err) - } } func testDatasetCreateService(t *testing.T) *DatasetService { @@ -144,6 +138,7 @@ func TestCreateDataset_ComponentParamsPopulated(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) insertCreateDatasetTenant(t, "tenant-1") + seedCreateDatasetDefaultChatModel(t, db, "tenant-1") ctx := t.Context() chunkMethod := "general" @@ -173,6 +168,24 @@ func TestCreateDataset_ComponentParamsPopulated(t *testing.T) { } } +func seedCreateDatasetDefaultChatModel(t *testing.T, db *gorm.DB, tenantID string) { + t.Helper() + if err := db.Model(&entity.UserTenant{}). + Where("user_id = ? AND tenant_id = ?", tenantID, tenantID). + Update("llm_id", "llm-default").Error; err != nil { + t.Fatalf("configure default chat model: %v", err) + } + for _, row := range []any{ + &entity.TenantModelProvider{ID: "provider-default-chat", UserID: tenantID, TenantID: tenantID, ProviderName: "OpenAI"}, + &entity.TenantModelInstance{ID: "instance-default-chat", ProviderID: "provider-default-chat", InstanceName: "default", Status: "active", Extra: "{}"}, + &entity.TenantModel{ID: "llm-default", ProviderID: "provider-default-chat", InstanceID: "instance-default-chat", ModelName: "gpt-4o", ModelType: int(entity.ModelTypeChat), Status: "active", Extra: "{}"}, + } { + if err := db.Create(row).Error; err != nil { + t.Fatalf("seed default chat model %T: %v", row, err) + } + } +} + // TestCreateDataset_ParentChildConfigReachesGeneralChunker verifies the public // parent-child setting controls the component that actually performs the // secondary split. Dropping this mapping silently leaves parent-child disabled. @@ -434,6 +447,47 @@ func TestCreateDataset_RejectsInvalidEmbeddingModel(t *testing.T) { } } +func TestCreateDataset_ExplicitEmbeddingDoesNotResolveInvalidDefault(t *testing.T) { + db := setupServiceTestDB(t) + pushServiceDB(t, db) + insertCreateDatasetTenant(t, "tenant-1") + if err := db.Model(&entity.UserTenant{}). + Where("user_id = ? AND tenant_id = ?", "tenant-1", "tenant-1"). + Updates(map[string]interface{}{"embd_id": "invalid-default", "tenant_embd_id": "missing-default"}).Error; err != nil { + t.Fatalf("configure invalid default embedding: %v", err) + } + if err := db.Create(&entity.TenantModelProvider{ + ID: "provider-explicit", UserID: "tenant-1", TenantID: "tenant-1", ProviderName: "OpenAI", + }).Error; err != nil { + t.Fatalf("create model provider: %v", err) + } + if err := db.Create(&entity.TenantModelInstance{ + ID: "instance-explicit", ProviderID: "provider-explicit", InstanceName: "default", Status: "active", Extra: "{}", + }).Error; err != nil { + t.Fatalf("create model instance: %v", err) + } + if err := db.Create(&entity.TenantModel{ + ID: "model-explicit", ProviderID: "provider-explicit", InstanceID: "instance-explicit", + ModelName: "embedding-test", ModelType: int(entity.ModelTypeEmbedding), Status: "active", Extra: "{}", + }).Error; err != nil { + t.Fatalf("create embedding model: %v", err) + } + + embeddingModel := "embedding-test@default@OpenAI" + result, code, err := testDatasetCreateService(t).CreateDataset(t.Context(), &service.CreateDatasetRequest{ + Name: "explicit-embedding", ParserID: sptr("general"), ParseType: iptr(1), EmbeddingModel: &embeddingModel, + }, testDatasetSubject("tenant-1")) + if err != nil { + t.Fatalf("CreateDataset with explicit embedding = %v, want success", err) + } + if code != common.CodeSuccess { + t.Fatalf("CreateDataset code = %d, want success", code) + } + if result["embedding_model"] != embeddingModel { + t.Fatalf("embedding_model = %#v, want %q", result["embedding_model"], embeddingModel) + } +} + func TestCreateDataset_SetsExplicitLanguage(t *testing.T) { db := setupServiceTestDB(t) pushServiceDB(t, db) diff --git a/internal/service/dataset/crud.go b/internal/service/dataset/crud.go index c9096e2f3c..21a046ec7e 100644 --- a/internal/service/dataset/crud.go +++ b/internal/service/dataset/crud.go @@ -229,12 +229,20 @@ func (d *DatasetService) CreateDataset(ctx context.Context, req *service.CreateD var parserConfigMap map[string]interface{} = parserConfig - embdID, tenantEmbdID, err := service.NewTenantService().GetDefaultModelRefsForUser(ctx, userID, tenantID, entity.ModelTypeEmbedding) - if err != nil { - return nil, common.CodeServerError, err + modelFactory := service.NewModelFactory() + identity := service.Identity{UserID: userID, TenantID: tenantID} + embdID, tenantEmbdID := "", "" + if embeddingModel == "" { + embeddingInfo, embeddingErr := modelFactory.ResolveDefaultInfo(ctx, identity, entity.ModelTypeEmbedding) + if embeddingErr == nil { + embdID = service.ModelReference(embeddingInfo) + tenantEmbdID = embeddingInfo.ID + } else if !errors.Is(embeddingErr, service.ErrNoDefaultModel) { + return nil, common.CodeServerError, embeddingErr + } } if embeddingModel != "" { - ok, message := d.verifyEmbeddingAvailability(ctx, embeddingModel, service.ModelAccess{UserID: userID, TenantID: tenantID}) + ok, message := d.verifyEmbeddingAvailability(ctx, embeddingModel, identity) if !ok { return nil, common.CodeDataError, errors.New(message) } @@ -242,7 +250,7 @@ func (d *DatasetService) CreateDataset(ctx context.Context, req *service.CreateD tenantEmbdID = "" } if embdID != "" && tenantEmbdID == "" { - target, err := service.NewModelFactory().ResolveInfo(ctx, service.ModelAccess{UserID: userID, TenantID: tenantID}, entity.ModelTypeEmbedding, embdID) + target, err := modelFactory.ResolveInfo(ctx, identity, entity.ModelTypeEmbedding, embdID) if err == nil { tenantEmbdID = target.ID } else { @@ -259,9 +267,12 @@ func (d *DatasetService) CreateDataset(ctx context.Context, req *service.CreateD } name = dedupedName - llmID, _, err := service.NewTenantService().GetDefaultModelRefsForUser(ctx, userID, tenantID, entity.ModelTypeChat) - if err != nil { - return nil, common.CodeServerError, err + llmID := "" + chatInfo, chatErr := modelFactory.ResolveDefaultInfo(ctx, identity, entity.ModelTypeChat) + if chatErr == nil { + llmID = service.ModelReference(chatInfo) + } else if !errors.Is(chatErr, service.ErrNoDefaultModel) { + return nil, common.CodeServerError, chatErr } parserConfig = service.ApplyComponentScopedParserConfig(parserConfig, llmID) diff --git a/internal/service/dataset/embedding.go b/internal/service/dataset/embedding.go index 18f0609cc1..ad59f4b124 100644 --- a/internal/service/dataset/embedding.go +++ b/internal/service/dataset/embedding.go @@ -62,7 +62,7 @@ func (d *DatasetService) CheckEmbedding(ctx context.Context, userID, datasetID s return nil, common.CodeServerError, errors.New("doc engine not initialized") } - embeddingModel, err := service.NewModelFactory().NewEmbeddingModel(ctx, service.ModelAccess{TenantID: kb.TenantID}, embeddingID) + embeddingModel, err := service.NewModelFactory().NewEmbeddingModel(ctx, service.Identity{TenantID: kb.TenantID}, embeddingID) if err != nil { return nil, common.CodeDataError, err } @@ -281,7 +281,7 @@ func (d *DatasetService) sampleRandomChunksWithVectors(ctx context.Context, tena return samples, nil } -func (d *DatasetService) verifyEmbeddingAvailability(ctx context.Context, embdID string, access service.ModelAccess) (bool, string) { +func (d *DatasetService) verifyEmbeddingAvailability(ctx context.Context, embdID string, access service.Identity) (bool, string) { _, err := service.NewModelFactory().ResolveInfo(ctx, access, entity.ModelTypeEmbedding, embdID) if err != nil { return false, err.Error() diff --git a/internal/service/dataset/metadata.go b/internal/service/dataset/metadata.go index cdfd8bc9ea..28b5600834 100644 --- a/internal/service/dataset/metadata.go +++ b/internal/service/dataset/metadata.go @@ -51,7 +51,10 @@ func (d *DatasetService) UpdateDocumentMetadataConfig(ctx context.Context, userI if _, tenantErr := d.tenantDAO.GetByID(ctx, dao.DB, kb.TenantID); tenantErr != nil { return nil, common.CodeServerError, errors.New("database operation failed") } - llmID, _, _ := service.NewTenantService().GetDefaultModelRefsForUser(ctx, userID, kb.TenantID, entity.ModelTypeChat) + llmID := "" + if modelInfo, modelErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{UserID: userID, TenantID: kb.TenantID}, entity.ModelTypeChat); modelErr == nil { + llmID = service.ModelReference(modelInfo) + } parserConfig = service.ApplyComponentScopedParserConfig(parserConfig, llmID) if err = d.documentDAO.UpdateByID(ctx, dao.DB, doc.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { @@ -159,7 +162,10 @@ func (d *DatasetService) UpdateMetadataConfig(ctx context.Context, datasetID, us } delete(parserConfig, "enable_metadata") delete(parserConfig, "built_in_metadata") - llmID, _, _ := service.NewTenantService().GetDefaultModelRefsForUser(ctx, userID, kb.TenantID, entity.ModelTypeChat) + llmID := "" + if modelInfo, modelErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{UserID: userID, TenantID: kb.TenantID}, entity.ModelTypeChat); modelErr == nil { + llmID = service.ModelReference(modelInfo) + } parserConfig = service.ApplyComponentScopedParserConfig(parserConfig, llmID) if err = d.kbDAO.UpdateByID(ctx, dao.DB, kb.ID, map[string]interface{}{"parser_config": parserConfig}); err != nil { diff --git a/internal/service/dataset/search.go b/internal/service/dataset/search.go index 0b2362a14d..8724332b5a 100644 --- a/internal/service/dataset/search.go +++ b/internal/service/dataset/search.go @@ -147,7 +147,7 @@ func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.Search common.Warn("Search service is not initialized for search_id", zap.String("searchID", searchID)) return nil, fmt.Errorf("invalid search_id") } - searchDetail, err := d.searchService.GetDetail(ctx, userID, searchID) + searchDetail, err := d.searchService.GetDetail(ctx, searchID) if err != nil { if ctx.Err() != nil { return nil, ctx.Err() @@ -215,7 +215,7 @@ func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.Search if metadataFilter != nil { method, _ := metadataFilter["method"].(string) if method == "auto" || method == "semi_auto" { - access := service.ModelAccess{TenantID: tenantIDs[0]} + access := service.Identity{TenantID: tenantIDs[0]} if chatID != "" { chatModelForFilter, err = modelFactory.NewChatModel(ctx, access, chatID) if err != nil { @@ -260,7 +260,7 @@ func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.Search } } if keyword { - chatModel, err := modelFactory.NewDefaultChatModel(ctx, service.ModelAccess{TenantID: tenantIDs[0]}) + chatModel, err := modelFactory.NewDefaultChatModel(ctx, service.Identity{TenantID: tenantIDs[0]}) if err != nil { common.Warn("Failed to get default chat model for LLM transformations", zap.Error(err)) } else { @@ -280,7 +280,7 @@ func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.Search // Determine embedding model var embeddingModel *modelModule.EmbeddingModel if kbRecords[0].EmbdID != "" { - embeddingModel, err = modelFactory.NewEmbeddingModel(ctx, service.ModelAccess{TenantID: tenantIDs[0]}, kbRecords[0].EmbdID) + embeddingModel, err = modelFactory.NewEmbeddingModel(ctx, service.Identity{TenantID: tenantIDs[0]}, kbRecords[0].EmbdID) if err != nil { return nil, fmt.Errorf("failed to get embedding model by embd_id: %w", err) } @@ -289,7 +289,7 @@ func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.Search // Get rerank model if rerankID is specified var rerankModel *modelModule.RerankModel if rerankID != "" { - rerankModel, err = modelFactory.NewRerankModel(ctx, service.ModelAccess{TenantID: tenantIDs[0]}, rerankID) + rerankModel, err = modelFactory.NewRerankModel(ctx, service.Identity{TenantID: tenantIDs[0]}, rerankID) if err != nil { return nil, fmt.Errorf("failed to get rerank model by rerank_id: %w", err) } diff --git a/internal/service/dataset/search_integration_test.go b/internal/service/dataset/search_integration_test.go index 7c44b1733f..8fa3355129 100644 --- a/internal/service/dataset/search_integration_test.go +++ b/internal/service/dataset/search_integration_test.go @@ -54,11 +54,11 @@ func TestSearchDatasetsSearchIDOverridesIntegration(t *testing.T) { } { t.Run(tt.name, func(t *testing.T) { searches := service.NewSearchService() - created, err := searches.CreateSearch(f.ctx, f.owner, "retrieval_"+tt.name, nil) + created, err := searches.CreateSearch(f.ctx, "retrieval_"+tt.name, nil) if err != nil { t.Fatal(err) } - _, err = searches.UpdateSearch(f.ctx, f.owner, created.SearchID, &service.UpdateSearchRequest{ + _, err = searches.UpdateSearch(f.ctx, created.SearchID, &service.UpdateSearchRequest{ Name: "retrieval_" + tt.name, SearchConfig: map[string]interface{}{ "keywords_similarity_weight": 1 - tt.weight, @@ -201,7 +201,7 @@ func newSearchIntegrationFixture(t *testing.T) *searchIntegrationFixture { } grantDatasetTestResource(t, owner, "kb", kb, owner, "owner") } - model, err := service.NewModelFactory().NewEmbeddingModel(ctx, service.ModelAccess{TenantID: owner}, info.ModelID+"@Builtin") + model, err := service.NewModelFactory().NewEmbeddingModel(ctx, service.Identity{TenantID: owner}, info.ModelID+"@Builtin") if err != nil { t.Fatal(err) } diff --git a/internal/service/dataset/update.go b/internal/service/dataset/update.go index 7a5486e6c6..909211fe61 100644 --- a/internal/service/dataset/update.go +++ b/internal/service/dataset/update.go @@ -271,14 +271,14 @@ func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, userID st } else { tenantEmbdID = "" } - modelAccess := service.ModelAccess{UserID: userID, TenantID: lockedKB.TenantID} - ok, message := d.verifyEmbeddingAvailability(ctx, effectiveEmbdID, modelAccess) + identity := service.Identity{UserID: userID, TenantID: lockedKB.TenantID} + ok, message := d.verifyEmbeddingAvailability(ctx, effectiveEmbdID, identity) if !ok { txCode = common.CodeDataError return errors.New(message) } if effectiveEmbdID != "" && tenantEmbdID == "" { - target, err := service.NewModelFactory().ResolveInfo(ctx, modelAccess, entity.ModelTypeEmbedding, effectiveEmbdID) + target, err := service.NewModelFactory().ResolveInfo(ctx, identity, entity.ModelTypeEmbedding, effectiveEmbdID) if err == nil { tenantEmbdID = target.ID } @@ -438,6 +438,9 @@ func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, userID st var membership entity.UserTenant if tenantErr := tx.Where("user_id = ? AND tenant_id = ? AND status = ?", userID, lockedKB.TenantID, entity.TenantMembershipStatusJoined).Take(&membership).Error; tenantErr == nil { llmID = membership.LLMID + if membership.TenantLLMID != nil && strings.TrimSpace(*membership.TenantLLMID) != "" { + llmID = strings.TrimSpace(*membership.TenantLLMID) + } } effectiveParserConfig = service.ApplyComponentScopedParserConfig( effectiveParserConfig, diff --git a/internal/service/dataset/update_test.go b/internal/service/dataset/update_test.go index 23473694e4..19ec23f2bc 100644 --- a/internal/service/dataset/update_test.go +++ b/internal/service/dataset/update_test.go @@ -1314,6 +1314,44 @@ func TestUpdateDataset_AcceptsValidComponentParams_Builtin(t *testing.T) { } } +func TestUpdateDatasetUsesTenantModelIDForExtractorDefault(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + if err := db.Model(&entity.UserTenant{}). + Where("user_id = ? AND tenant_id = ?", "tenant-1", "tenant-1"). + Updates(map[string]interface{}{ + "llm_id": "glm-4-flash@CI@ZHIPU-AI", + "tenant_llm_id": "chat-model-id", + }).Error; err != nil { + t.Fatalf("set default chat model: %v", err) + } + + ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1") + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ + ParserConfig: map[string]interface{}{ + "Parser:HipSignsRhyme": map[string]interface{}{ + "pdf": map[string]interface{}{"parse_method": "deepdoc"}, + }, + }, + }) + if err != nil || code != common.CodeSuccess { + t.Fatalf("UpdateDataset failed: code=%d err=%v", code, err) + } + + persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1") + if err != nil { + t.Fatalf("get updated dataset: %v", err) + } + extractor, ok := persisted.ParserConfig["Extractor:AutoExtractDefault"].(map[string]interface{}) + if !ok { + t.Fatalf("extractor parser config = %#v", persisted.ParserConfig["Extractor:AutoExtractDefault"]) + } + if extractor["llm_id"] != "chat-model-id" { + t.Fatalf("extractor llm_id = %#v, want tenant model ID", extractor["llm_id"]) + } +} + // extractorNodeMetadataInTest returns the metadata object stored on the first // Extractor node of a parser_config (the component-scoped form). func extractorNodeMetadataInTest(t *testing.T, cfg map[string]interface{}) map[string]interface{} { diff --git a/internal/service/document/document_applykb_chain_test.go b/internal/service/document/document_applykb_chain_test.go index a3d8d3200c..664289761d 100644 --- a/internal/service/document/document_applykb_chain_test.go +++ b/internal/service/document/document_applykb_chain_test.go @@ -308,7 +308,7 @@ func TestIngest_ApplyKB_RealCreatedDatasetChain(t *testing.T) { // resolved embedding model id, keeping CreateDataset on its real path // (no embedding override). if err := dao.DB.Model(&entity.UserTenant{}).Where("user_id = ? AND tenant_id = ?", "user-1", "tenant-1").Update( - "tenant_embd_id", "embd-1", + "tenant_embd_id", "embd-default", ).Error; err != nil { t.Fatalf("set membership embd id: %v", err) } diff --git a/internal/service/document/document_dataset_update.go b/internal/service/document/document_dataset_update.go index 653f71f662..1f3fc4074d 100644 --- a/internal/service/document/document_dataset_update.go +++ b/internal/service/document/document_dataset_update.go @@ -292,7 +292,10 @@ func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, dat } } pipelinepkg.ApplyParentChildChunkerConfig(cleaned, parentChildOverrides) - llmID, _, _ := service.NewTenantService().GetDefaultModelRefsForUser(ctx, userID, kb.TenantID, entity.ModelTypeChat) + llmID := "" + if modelInfo, modelErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{UserID: userID, TenantID: kb.TenantID}, entity.ModelTypeChat); modelErr == nil { + llmID = service.ModelReference(modelInfo) + } cleaned = service.ApplyComponentScopedParserConfig(cleaned, llmID) parserUpdates["parser_config"] = cleaned } @@ -320,11 +323,12 @@ func (s *DocumentService) UpdateDatasetDocument(ctx context.Context, userID, dat common.Warn("document pipeline/version switch: failed to resolve component params defaults, keeping prior parser_config", zap.String("parserID", effParserID), zap.Error(cpErr)) } else { - llmID, _, llmErr := service.NewTenantService().GetDefaultModelRefsForUser(ctx, userID, kb.TenantID, entity.ModelTypeChat) - if llmErr != nil { + llmID := "" + if modelInfo, llmErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{UserID: userID, TenantID: kb.TenantID}, entity.ModelTypeChat); llmErr == nil { + llmID = service.ModelReference(modelInfo) + } else if !errors.Is(llmErr, service.ErrNoDefaultModel) { common.Warn("document pipeline/version switch: failed to resolve default chat model, falling back to no llm binding", zap.String("tenant_id", kb.TenantID), zap.Error(llmErr)) - llmID = "" } cleaned := service.ApplyComponentScopedParserConfig(cpDefaults, llmID) parserUpdates["parser_config"] = cleaned diff --git a/internal/service/document/document_parser_selection.go b/internal/service/document/document_parser_selection.go index 22cc347112..d75f7f0295 100644 --- a/internal/service/document/document_parser_selection.go +++ b/internal/service/document/document_parser_selection.go @@ -201,7 +201,10 @@ func rebuildParserConfig(ctx context.Context, kb *entity.Knowledgebase, target, cleaned["metadata"] = meta } pipelinepkg.ApplyParentChildChunkerConfig(cleaned, raw) - llmID, _ := service.NewTenantService().GetDefaultModelNameForTenant(ctx, kb.TenantID, entity.ModelTypeChat) + llmID := "" + if modelInfo, modelErr := service.NewModelFactory().ResolveDefaultInfo(ctx, service.Identity{TenantID: kb.TenantID}, entity.ModelTypeChat); modelErr == nil { + llmID = service.ModelReference(modelInfo) + } cleaned = service.ApplyComponentScopedParserConfig(cleaned, llmID) return cleaned, true } diff --git a/internal/service/document/document_parser_selection_test.go b/internal/service/document/document_parser_selection_test.go index 0c6454a499..b7be44a647 100644 --- a/internal/service/document/document_parser_selection_test.go +++ b/internal/service/document/document_parser_selection_test.go @@ -8,7 +8,6 @@ import ( "ragflow/internal/dao" "ragflow/internal/entity" - "ragflow/internal/permission/testsupport" "ragflow/internal/service" "ragflow/internal/storage" syncerconnector "ragflow/internal/syncer/connector" @@ -150,19 +149,7 @@ func tableParserConfig() entity.JSONMap { func seedDocumentParserTenant(t *testing.T, tenantID, ownerID string) { t.Helper() - if err := dao.DB.Create(&entity.Tenant{ID: tenantID}).Error; err != nil { - t.Fatalf("insert tenant: %v", err) - } - status := string(entity.StatusValid) - if err := dao.DB.Create(&entity.UserTenant{ - ID: "owner-" + tenantID, UserID: ownerID, TenantID: tenantID, Role: "owner", Status: &status, - LLMID: "llm-default", EmbdID: "embd-default", ASRID: "asr-default", - }).Error; err != nil { - t.Fatalf("insert tenant owner membership: %v", err) - } - if err := testsupport.PrepareEnterprisePermissionFixture(dao.DB, tenantID, ownerID); err != nil { - t.Fatalf("prepare parser tenant permission fixture: %v", err) - } + insertUserTenantForAccessCheck(t, ownerID, tenantID) } // TestResolveDocumentParser_Priority pins the resolution order: the dataset's diff --git a/internal/service/document/document_test.go b/internal/service/document/document_test.go index 040af5f923..5fb1e1ef54 100644 --- a/internal/service/document/document_test.go +++ b/internal/service/document/document_test.go @@ -549,6 +549,9 @@ func setupServiceTestDBWithDSN(t *testing.T, dsn string) *gorm.DB { &entity.User{}, &entity.Tenant{}, &entity.UserTenant{}, + &entity.TenantModelProvider{}, + &entity.TenantModelInstance{}, + &entity.TenantModel{}, &entity.UserCanvas{}, &entity.API4Conversation{}, } @@ -1493,6 +1496,15 @@ func TestUploadEmptyDocument_CreatesVirtualDocumentAndFileLink(t *testing.T) { func insertUserTenantForAccessCheck(t *testing.T, userID, tenantID string) { t.Helper() + modelID := func(defaultID string) string { + if userID == "user-1" && tenantID == "tenant-1" { + return defaultID + } + return fmt.Sprintf("%s-%s-%s", defaultID, userID, tenantID) + } + chatModelID := modelID("llm-default") + embeddingModelID := modelID("embd-default") + asrModelID := modelID("asr-default") // Insert user if not exists (email is NOT NULL, password is nullable pointer) var existingUser entity.User if err := dao.DB.Where("id = ?", userID).First(&existingUser).Error; err != nil { @@ -1518,9 +1530,9 @@ func insertUserTenantForAccessCheck(t *testing.T, userID, tenantID string) { TenantID: tenantID, Role: "owner", Status: sptr(entity.TenantMembershipStatusJoined), - LLMID: "llm-default", - EmbdID: "embd-default", - ASRID: "asr-default", + LLMID: chatModelID, + EmbdID: embeddingModelID, + ASRID: asrModelID, } if err = dao.DB.Create(ut).Error; err != nil { t.Fatalf("insert test user_tenant: %v", err) @@ -1528,9 +1540,28 @@ func insertUserTenantForAccessCheck(t *testing.T, userID, tenantID string) { } if err := dao.DB.Model(&entity.UserTenant{}). Where("user_id = ? AND tenant_id = ?", userID, tenantID). - Updates(map[string]interface{}{"llm_id": "llm-default", "embd_id": "embd-default", "asr_id": "asr-default"}).Error; err != nil { + Updates(map[string]interface{}{"llm_id": chatModelID, "embd_id": embeddingModelID, "asr_id": asrModelID}).Error; err != nil { t.Fatalf("set test owner model defaults: %v", err) } + providerID := fmt.Sprintf("test-default-provider-%s-%s", userID, tenantID) + instanceID := fmt.Sprintf("test-default-instance-%s-%s", userID, tenantID) + provider := &entity.TenantModelProvider{ID: providerID, UserID: userID, TenantID: tenantID, ProviderName: "OpenAI"} + if err := dao.DB.Where("id = ?", providerID).FirstOrCreate(provider).Error; err != nil { + t.Fatalf("seed test default provider: %v", err) + } + instance := &entity.TenantModelInstance{ID: instanceID, ProviderID: providerID, InstanceName: "default", Status: "active", Extra: "{}"} + if err := dao.DB.Where("id = ?", instanceID).FirstOrCreate(instance).Error; err != nil { + t.Fatalf("seed test default instance: %v", err) + } + for _, model := range []*entity.TenantModel{ + {ID: chatModelID, ProviderID: providerID, InstanceID: instanceID, ModelName: "gpt-4o", ModelType: int(entity.ModelTypeChat), Status: "active", Extra: "{}"}, + {ID: embeddingModelID, ProviderID: providerID, InstanceID: instanceID, ModelName: "text-embedding-3-small", ModelType: int(entity.ModelTypeEmbedding), Status: "active", Extra: "{}"}, + {ID: asrModelID, ProviderID: providerID, InstanceID: instanceID, ModelName: "whisper-1", ModelType: int(entity.ModelTypeSpeech2Text), Status: "active", Extra: "{}"}, + } { + if err := dao.DB.Where("id = ?", model.ID).FirstOrCreate(model).Error; err != nil { + t.Fatalf("seed test default model %s: %v", model.ID, err) + } + } if err := testsupport.PrepareEnterprisePermissionFixture(dao.DB, tenantID, userID); err != nil { t.Fatalf("prepare document permission fixture: %v", err) } diff --git a/internal/service/generator.go b/internal/service/generator.go index 9f7d570d99..6fd9d5847b 100644 --- a/internal/service/generator.go +++ b/internal/service/generator.go @@ -126,7 +126,7 @@ func CrossLanguages(ctx context.Context, tenantID string, llmID string, query st var chatModel *modelModule.ChatModel var err error - access := ModelAccess{TenantID: tenantID} + access := Identity{TenantID: tenantID} if llmID != "" { chatModel, err = modelFactory.NewChatModel(ctx, access, llmID) if err != nil { diff --git a/internal/service/identity.go b/internal/service/identity.go new file mode 100644 index 0000000000..cde34666c2 --- /dev/null +++ b/internal/service/identity.go @@ -0,0 +1,47 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "context" + "strings" + + "ragflow/internal/permission" +) + +// Identity identifies a user in a tenant. TenantID may be empty when an +// edition derives the effective tenant from the resource or user. +type Identity struct { + UserID string + TenantID string +} + +// ResolveIdentity returns the authenticated user and their active tenant. +// Explicit Identity values remain available to background work without a +// request context. +func ResolveIdentity(ctx context.Context) (Identity, error) { + subject, ok := permission.SubjectFromContext(ctx) + if !ok || strings.TrimSpace(subject.UserID) == "" { + return Identity{}, permission.ErrUnauthenticated + } + + tenantID, err := requireCurrentTenantID(ctx, subject.UserID) + if err != nil { + return Identity{}, err + } + + return Identity{UserID: strings.TrimSpace(subject.UserID), TenantID: tenantID}, nil +} diff --git a/internal/service/memory.go b/internal/service/memory.go index 30a6e0ade5..cf60e9605e 100644 --- a/internal/service/memory.go +++ b/internal/service/memory.go @@ -503,7 +503,7 @@ func (s *MemoryService) CreateMemory(ctx context.Context, userID string, req *Cr // tenant_model row) — we leave the tenant_*_id fields nil and proceed. modelFactory := NewModelFactory() if req.LLMID != "" && req.TenantLLMID == nil { - info, err := modelFactory.ResolveInfo(ctx, ModelAccess{UserID: userID, TenantID: tenantID}, entity.ModelTypeChat, req.LLMID) + info, err := modelFactory.ResolveInfo(ctx, Identity{UserID: userID, TenantID: tenantID}, entity.ModelTypeChat, req.LLMID) if err != nil { common.Warn("CreateMemory: failed to resolve tenant LLM id", zap.String("tenant_id", tenantID), zap.String("llm_id", req.LLMID), zap.Error(err)) @@ -513,7 +513,7 @@ func (s *MemoryService) CreateMemory(ctx context.Context, userID string, req *Cr } } if req.EmbdID != "" && req.TenantEmbdID == nil { - info, err := modelFactory.ResolveInfo(ctx, ModelAccess{UserID: userID, TenantID: tenantID}, entity.ModelTypeEmbedding, req.EmbdID) + info, err := modelFactory.ResolveInfo(ctx, Identity{UserID: userID, TenantID: tenantID}, entity.ModelTypeEmbedding, req.EmbdID) if err != nil { common.Warn("CreateMemory: failed to resolve tenant embedding id", zap.String("tenant_id", tenantID), zap.String("embd_id", req.EmbdID), zap.Error(err)) @@ -673,7 +673,7 @@ func (s *MemoryService) UpdateMemory(ctx context.Context, userID string, memoryI if req.LLMID != nil { updateDict["llm_id"] = *req.LLMID if req.TenantLLMID == nil && *req.LLMID != "" { - info, err := modelFactory.ResolveInfo(ctx, ModelAccess{TenantID: ownerTenantID}, entity.ModelTypeChat, *req.LLMID) + info, err := modelFactory.ResolveInfo(ctx, Identity{TenantID: ownerTenantID}, entity.ModelTypeChat, *req.LLMID) if err != nil { common.Warn("UpdateMemory: failed to resolve tenant LLM id", zap.String("tenant_id", ownerTenantID), zap.String("llm_id", *req.LLMID), zap.Error(err)) @@ -686,7 +686,7 @@ func (s *MemoryService) UpdateMemory(ctx context.Context, userID string, memoryI if req.EmbdID != nil { updateDict["embd_id"] = *req.EmbdID if req.TenantEmbdID == nil && *req.EmbdID != "" { - info, err := modelFactory.ResolveInfo(ctx, ModelAccess{TenantID: ownerTenantID}, entity.ModelTypeEmbedding, *req.EmbdID) + info, err := modelFactory.ResolveInfo(ctx, Identity{TenantID: ownerTenantID}, entity.ModelTypeEmbedding, *req.EmbdID) if err != nil { common.Warn("UpdateMemory: failed to resolve tenant embedding id", zap.String("tenant_id", ownerTenantID), zap.String("embd_id", *req.EmbdID), zap.Error(err)) @@ -1539,7 +1539,7 @@ func memoryFusionWeights(keywordsSimilarityWeight float64) string { } func (s *MemoryService) memoryMessageDenseExpr(ctx context.Context, question string, memory *entity.Memory, topN int, similarityThreshold float64) (*enginetypes.MatchDenseExpr, error) { - embeddingModel, err := NewModelFactory().NewEmbeddingModel(ctx, ModelAccess{TenantID: memory.TenantID}, memory.EmbdID) + embeddingModel, err := NewModelFactory().NewEmbeddingModel(ctx, Identity{TenantID: memory.TenantID}, memory.EmbdID) if err != nil { return nil, err } diff --git a/internal/service/memory_extractor.go b/internal/service/memory_extractor.go index d1fe1ab58d..7b98d689d6 100644 --- a/internal/service/memory_extractor.go +++ b/internal/service/memory_extractor.go @@ -496,7 +496,7 @@ func (s *MemoryMessageService) extractByLLM(ctx context.Context, mem *CreateMemo if mem.TenantLLMID != nil && *mem.TenantLLMID != "" { llmRef = *mem.TenantLLMID } - chatModel, err := NewModelFactory().NewChatModel(ctx, ModelAccess{TenantID: mem.TenantID}, llmRef) + chatModel, err := NewModelFactory().NewChatModel(ctx, Identity{TenantID: mem.TenantID}, llmRef) if err != nil { return nil, fmt.Errorf("resolve chat model: %w", classifyMemoryTaskDependencyError(err)) } diff --git a/internal/service/memory_message_service.go b/internal/service/memory_message_service.go index f309107e2c..990501abc5 100644 --- a/internal/service/memory_message_service.go +++ b/internal/service/memory_message_service.go @@ -355,7 +355,7 @@ func (s *MemoryMessageService) embedAndSaveMessages(ctx context.Context, mem *Cr for i, message := range messages { contents[i], _ = message["content"].(string) } - embeddingModel, err := NewModelFactory().NewEmbeddingModel(ctx, ModelAccess{TenantID: mem.TenantID}, mem.EmbdID) + embeddingModel, err := NewModelFactory().NewEmbeddingModel(ctx, Identity{TenantID: mem.TenantID}, mem.EmbdID) if err != nil { return err } diff --git a/internal/service/memory_message_test.go b/internal/service/memory_message_test.go index 277e03d124..96aa2c5fd8 100644 --- a/internal/service/memory_message_test.go +++ b/internal/service/memory_message_test.go @@ -467,14 +467,16 @@ func TestUpdateMemoryTeamMemberResolvesModelsAgainstOwnerTenant(t *testing.T) { for _, row := range []struct { providerID string + userID string tenantID string modelID string }{ - {providerID: "provider-owner", tenantID: "owner-1", modelID: "tenant-llm-owner"}, - {providerID: "provider-member", tenantID: "member-1", modelID: "tenant-llm-member"}, + {providerID: "provider-owner", userID: "owner-1", tenantID: "owner-1", modelID: "tenant-llm-owner"}, + {providerID: "provider-member", userID: "member-1", tenantID: "member-1", modelID: "tenant-llm-member"}, } { if err := dao.DB.Create(&entity.TenantModelProvider{ ID: row.providerID, + UserID: row.userID, ProviderName: "OpenAI", TenantID: row.tenantID, }).Error; err != nil { diff --git a/internal/service/model_call.go b/internal/service/model_call.go index c62de77b93..65de88f17d 100644 --- a/internal/service/model_call.go +++ b/internal/service/model_call.go @@ -22,10 +22,7 @@ import ( "fmt" "strings" - "gorm.io/gorm" - "ragflow/internal/common" - "ragflow/internal/dao" modelModule "ragflow/internal/entity/models" ) @@ -34,38 +31,21 @@ import ( // runtime configuration. Provider, instance, API configuration, and driver // resolution are handled internally. type ModelCallService struct { - providerService *ModelProviderService - modelFactory *ModelFactory + modelFactory *ModelFactory } -// NewModelCallService creates a model-call service with the standard model -// provider and model resolver. +// NewModelCallService creates a model-call service with the standard resolver. func NewModelCallService() *ModelCallService { - return NewModelCallServiceWithProviderService(NewModelProviderService(), NewModelFactory()) -} - -// NewModelCallServiceWithProviderService creates a model-call service using -// an existing provider service for provider metadata and balance operations. -func NewModelCallServiceWithProviderService(providerService *ModelProviderService, modelFactory *ModelFactory) *ModelCallService { - if providerService == nil { - providerService = NewModelProviderService() - } - if modelFactory == nil { - modelFactory = NewModelFactory() - } - return &ModelCallService{ - providerService: providerService, - modelFactory: modelFactory, - } + return &ModelCallService{modelFactory: NewModelFactory()} } // ChatToModelWithMessages sends messages to the model selected by modelRef. -func (s *ModelCallService) ChatToModelWithMessages(ctx context.Context, modelRef, userID string, messages []modelModule.Message, config *modelModule.ChatConfig, usage *common.ModelUsage) (*modelModule.ChatResponse, common.ErrorCode, error) { - access, tenantID, code, err := s.modelAccessForUser(ctx, userID) +func (s *ModelCallService) ChatToModelWithMessages(ctx context.Context, modelRef string, messages []modelModule.Message, config *modelModule.ChatConfig, usage *common.ModelUsage) (*modelModule.ChatResponse, common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) if err != nil { return nil, code, err } - chatModel, err := s.modelFactory.NewChatModel(ctx, access, modelRef) + chatModel, err := s.modelFactory.NewChatModel(ctx, identity, modelRef) if err != nil { return nil, common.CodeNotFound, err } @@ -73,7 +53,7 @@ func (s *ModelCallService) ChatToModelWithMessages(ctx context.Context, modelRef config = &modelModule.ChatConfig{} } info := chatModel.Info() - populateModelUsage(usage, info, chatModel.APIConfig, tenantID, userID) + populateModelUsage(usage, info, chatModel.APIConfig, identity.TenantID, identity.UserID) response, err := chatModel.ChatWithMessages(ctx, messages, config, usage) if err != nil { @@ -84,31 +64,31 @@ func (s *ModelCallService) ChatToModelWithMessages(ctx context.Context, modelRef } return response, common.CodeSuccess, nil } -func (s *ModelCallService) ChatToModelStreamWithSender(ctx context.Context, modelRef, userID string, messages []modelModule.Message, config *modelModule.ChatConfig, usage *common.ModelUsage, sender func(*string, *string) error) (common.ErrorCode, error) { - access, tenantID, code, err := s.modelAccessForUser(ctx, userID) +func (s *ModelCallService) ChatToModelStreamWithSender(ctx context.Context, modelRef string, messages []modelModule.Message, config *modelModule.ChatConfig, usage *common.ModelUsage, sender func(*string, *string) error) (common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) if err != nil { return code, err } - chatModel, err := s.modelFactory.NewChatModel(ctx, access, modelRef) + chatModel, err := s.modelFactory.NewChatModel(ctx, identity, modelRef) if err != nil { return common.CodeNotFound, err } if config == nil { config = &modelModule.ChatConfig{} } - populateModelUsage(usage, chatModel.Info(), chatModel.APIConfig, tenantID, userID) + populateModelUsage(usage, chatModel.Info(), chatModel.APIConfig, identity.TenantID, identity.UserID) if err := chatModel.ChatStreamlyWithSender(ctx, messages, config, usage, sender); err != nil { return common.CodeServerError, err } return common.CodeSuccess, nil } -func (s *ModelCallService) EmbedText(ctx context.Context, modelRef, userID string, texts []string, config *modelModule.EmbeddingConfig) ([]modelModule.EmbeddingData, common.ErrorCode, error) { - access, _, code, err := s.modelAccessForUser(ctx, userID) +func (s *ModelCallService) EmbedText(ctx context.Context, modelRef string, texts []string, config *modelModule.EmbeddingConfig) ([]modelModule.EmbeddingData, common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) if err != nil { return nil, code, err } - embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, access, modelRef) + embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, identity, modelRef) if err != nil { return nil, common.CodeNotFound, err } @@ -132,12 +112,12 @@ func (s *ModelCallService) EmbedText(ctx context.Context, modelRef, userID strin } return embeddings, common.CodeSuccess, nil } -func (s *ModelCallService) RerankDocument(ctx context.Context, modelRef, userID, query string, documents []string, config *modelModule.RerankConfig) (*modelModule.RerankResponse, common.ErrorCode, error) { - access, _, code, err := s.modelAccessForUser(ctx, userID) +func (s *ModelCallService) RerankDocument(ctx context.Context, modelRef, query string, documents []string, config *modelModule.RerankConfig) (*modelModule.RerankResponse, common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) if err != nil { return nil, code, err } - rerankModel, err := s.modelFactory.NewRerankModel(ctx, access, modelRef) + rerankModel, err := s.modelFactory.NewRerankModel(ctx, identity, modelRef) if err != nil { return nil, common.CodeNotFound, err } @@ -154,12 +134,12 @@ func (s *ModelCallService) RerankDocument(ctx context.Context, modelRef, userID, } return response, common.CodeSuccess, nil } -func (s *ModelCallService) TranscribeAudio(ctx context.Context, modelRef, userID string, audioFile *string, config *modelModule.ASRConfig) (*modelModule.ASRResponse, common.ErrorCode, error) { - access, _, code, err := s.modelAccessForUser(ctx, userID) +func (s *ModelCallService) TranscribeAudio(ctx context.Context, modelRef string, audioFile *string, config *modelModule.ASRConfig) (*modelModule.ASRResponse, common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) if err != nil { return nil, code, err } - asrModel, err := s.modelFactory.NewASRModel(ctx, access, modelRef) + asrModel, err := s.modelFactory.NewASRModel(ctx, identity, modelRef) if err != nil { return nil, common.CodeNotFound, err } @@ -176,12 +156,12 @@ func (s *ModelCallService) TranscribeAudio(ctx context.Context, modelRef, userID } return response, common.CodeSuccess, nil } -func (s *ModelCallService) TranscribeAudioStream(ctx context.Context, modelRef, userID string, audioFile *string, config *modelModule.ASRConfig, sender func(*string, *string) error) (common.ErrorCode, error) { - access, _, code, err := s.modelAccessForUser(ctx, userID) +func (s *ModelCallService) TranscribeAudioStream(ctx context.Context, modelRef string, audioFile *string, config *modelModule.ASRConfig, sender func(*string, *string) error) (common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) if err != nil { return code, err } - asrModel, err := s.modelFactory.NewASRModel(ctx, access, modelRef) + asrModel, err := s.modelFactory.NewASRModel(ctx, identity, modelRef) if err != nil { return common.CodeNotFound, err } @@ -193,28 +173,18 @@ func (s *ModelCallService) TranscribeAudioStream(ctx context.Context, modelRef, } return common.CodeSuccess, nil } -func (s *ModelCallService) AudioSpeech(ctx context.Context, modelRef, userID string, audioContent *string, config *modelModule.TTSConfig) (*modelModule.TTSResponse, common.ErrorCode, error) { - if s == nil || s.modelFactory == nil { - return nil, common.CodeServerError, errors.New("model call service is not initialized") - } - tenantID, tenantErr := s.ownerTenantID(ctx, userID) - access := ModelAccess{UserID: userID, TenantID: tenantID} - if tenantErr != nil { - if strings.TrimSpace(modelRef) != "" { - return nil, common.CodeNotFound, tenantErr - } - // Some internal audio dispatchers pass a tenant ID in userID. - access = ModelAccess{TenantID: userID} - tenantID = userID +func (s *ModelCallService) AudioSpeech(ctx context.Context, modelRef string, audioContent *string, config *modelModule.TTSConfig) (*modelModule.TTSResponse, common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) + if err != nil { + return nil, code, err } var ( ttsModel *modelModule.TTSModel - err error ) if strings.TrimSpace(modelRef) == "" { - ttsModel, err = s.modelFactory.NewDefaultTTSModel(ctx, access) + ttsModel, err = s.modelFactory.NewDefaultTTSModel(ctx, identity) } else { - ttsModel, err = s.modelFactory.NewTTSModel(ctx, access, modelRef) + ttsModel, err = s.modelFactory.NewTTSModel(ctx, identity, modelRef) } if err != nil { return nil, common.CodeNotFound, err @@ -236,7 +206,7 @@ func (s *ModelCallService) AudioSpeechForTenant(ctx context.Context, modelRef, t if s == nil || s.modelFactory == nil { return nil, common.CodeServerError, errors.New("model call service is not initialized") } - access := ModelAccess{TenantID: tenantID} + access := Identity{TenantID: tenantID} var ( ttsModel *modelModule.TTSModel err error @@ -262,12 +232,12 @@ func (s *ModelCallService) AudioSpeechForTenant(ctx context.Context, modelRef, t } return response, common.CodeSuccess, nil } -func (s *ModelCallService) AudioSpeechStream(ctx context.Context, modelRef, userID string, audioContent *string, config *modelModule.TTSConfig, sender func(*string, *string) error) (common.ErrorCode, error) { - access, _, code, err := s.modelAccessForUser(ctx, userID) +func (s *ModelCallService) AudioSpeechStream(ctx context.Context, modelRef string, audioContent *string, config *modelModule.TTSConfig, sender func(*string, *string) error) (common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) if err != nil { return code, err } - ttsModel, err := s.modelFactory.NewTTSModel(ctx, access, modelRef) + ttsModel, err := s.modelFactory.NewTTSModel(ctx, identity, modelRef) if err != nil { return common.CodeNotFound, err } @@ -279,12 +249,12 @@ func (s *ModelCallService) AudioSpeechStream(ctx context.Context, modelRef, user } return common.CodeSuccess, nil } -func (s *ModelCallService) OCRFile(ctx context.Context, modelRef, userID string, content []byte, url *string, config *modelModule.OCRConfig) (*modelModule.OCRFileResponse, common.ErrorCode, error) { - access, _, code, err := s.modelAccessForUser(ctx, userID) +func (s *ModelCallService) OCRFile(ctx context.Context, modelRef string, content []byte, url *string, config *modelModule.OCRConfig) (*modelModule.OCRFileResponse, common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) if err != nil { return nil, code, err } - ocrModel, err := s.modelFactory.NewOCRModel(ctx, access, modelRef) + ocrModel, err := s.modelFactory.NewOCRModel(ctx, identity, modelRef) if err != nil { return nil, common.CodeNotFound, err } @@ -303,46 +273,24 @@ func (s *ModelCallService) OCRFile(ctx context.Context, modelRef, userID string, } // ParseFile parses content or url with the document parser selected by -// modelRef. doc_parse is a provider catalog capability rather than a member -// of entity.ModelType, so this method uses the existing model-info resolver. -func (s *ModelCallService) ParseFile(ctx context.Context, modelRef, userID string, content []byte, url *string, config *modelModule.ParseFileConfig) (*modelModule.ParseFileResponse, common.ErrorCode, error) { - info, err := s.resolveModelInfo(ctx, modelRef, userID) +// modelRef. A zero model type resolves catalog-only capabilities such as doc_parse. +func (s *ModelCallService) ParseFile(ctx context.Context, modelRef string, content []byte, url *string, config *modelModule.ParseFileConfig) (*modelModule.ParseFileResponse, common.ErrorCode, error) { + identity, code, err := s.identityForRequest(ctx) + if err != nil { + return nil, code, err + } + target, err := s.modelFactory.resolveConfig(ctx, identity, 0, modelRef) if err != nil { return nil, common.CodeNotFound, err } - if info.ModelInfo == nil || !info.ModelInfo.ModelTypeMap["doc_parse"] { + if target.ModelInfo == nil || !target.ModelInfo.ModelTypeMap["doc_parse"] { return nil, common.CodeNotFound, fmt.Errorf("model %q is not a document parser", modelRef) } if config == nil { config = &modelModule.ParseFileConfig{} } - modelDriver := info.ProviderInfo.ModelDriver - if info.ModelEntity != nil { - if info.ModelEntity.Status != "active" { - return nil, common.CodeNotFound, errors.New("model is inactive") - } - region := "" - baseURL := "" - if info.APIConfig != nil { - if info.APIConfig.Region != nil { - region = *info.APIConfig.Region - } - if info.APIConfig.BaseURL != nil { - baseURL = *info.APIConfig.BaseURL - } - } - modelDriver, err = newModelDriverForBaseURL(info.ProviderInfo.ModelDriver, info.ProviderEntity.ProviderName, region, baseURL) - if err != nil { - return nil, common.CodeServerError, err - } - } - - modelName := info.ModelInfo.Name - if info.ModelEntity != nil && info.ModelEntity.ModelName != "" { - modelName = info.ModelEntity.ModelName - } - response, err := modelDriver.ParseFile(ctx, &modelName, content, url, info.APIConfig, config, nil) + response, err := target.Driver.ParseFile(ctx, &target.ModelName, content, url, target.APIConfig, config, nil) if err != nil { return nil, common.CodeServerError, err } @@ -352,15 +300,15 @@ func (s *ModelCallService) ParseFile(ctx context.Context, modelRef, userID strin return response, common.CodeSuccess, nil } -func (s *ModelCallService) modelAccessForUser(ctx context.Context, userID string) (ModelAccess, string, common.ErrorCode, error) { - if s == nil || s.providerService == nil || s.modelFactory == nil { - return ModelAccess{}, "", common.CodeServerError, errors.New("model call service is not initialized") +func (s *ModelCallService) identityForRequest(ctx context.Context) (Identity, common.ErrorCode, error) { + if s == nil || s.modelFactory == nil { + return Identity{}, common.CodeServerError, errors.New("model call service is not initialized") } - tenantID, err := s.ownerTenantID(ctx, userID) + identity, err := ResolveIdentity(ctx) if err != nil { - return ModelAccess{}, "", common.CodeNotFound, err + return Identity{}, common.CodeNotFound, err } - return ModelAccess{UserID: userID, TenantID: tenantID}, tenantID, common.CodeSuccess, nil + return identity, common.CodeSuccess, nil } func populateModelUsage(usage *common.ModelUsage, info *modelModule.ModelInfo, apiConfig *modelModule.APIConfig, tenantID, userID string) { @@ -376,31 +324,3 @@ func populateModelUsage(usage *common.ModelUsage, info *modelModule.ModelInfo, a usage.APIKey = *apiConfig.ApiKey } } - -func (s *ModelCallService) ownerTenantID(ctx context.Context, userID string) (string, error) { - return requireCurrentTenantID(ctx, userID) -} - -func (s *ModelCallService) resolveModelInfo(ctx context.Context, modelRef, userID string) (*ModelInstanceAndProviderInfo, error) { - if s == nil || s.providerService == nil { - return nil, errors.New("model call service is not initialized") - } - modelRef = strings.TrimSpace(modelRef) - if modelRef == "" { - return nil, errors.New("model ref is required") - } - - model, err := s.providerService.modelDAO.GetByID(ctx, dao.DB, modelRef) - if err == nil && model != nil { - return s.providerService.getModelInstanceAndProviderByID(ctx, &modelRef, userID, nil) - } - if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { - return nil, err - } - - modelName, instanceName, providerName, err := common.ParseCompositeModelName(modelRef) - if err != nil { - return nil, err - } - return s.providerService.getModelInstanceAndProviderByName(ctx, &providerName, &instanceName, &modelName, userID, nil) -} diff --git a/internal/service/model_chat.go b/internal/service/model_chat.go index 896c7dfd3a..1f351153fc 100644 --- a/internal/service/model_chat.go +++ b/internal/service/model_chat.go @@ -69,7 +69,7 @@ func chatStreamWithContext(ctx context.Context, chatModel *modelModule.ChatModel // channel-based streaming interface. type TenantStreamAdapter struct { Factory *ModelFactory - TenantID string + Identity Identity ModelID string } @@ -77,7 +77,7 @@ func (a *TenantStreamAdapter) ChatStream(ctx context.Context, messages []modelMo if a.Factory == nil { return nil, nil, fmt.Errorf("model factory not configured") } - chatModel, err := a.Factory.NewChatModel(ctx, ModelAccess{TenantID: a.TenantID}, a.ModelID) + chatModel, err := a.Factory.NewChatModel(ctx, a.Identity, a.ModelID) if err != nil { return nil, nil, err } diff --git a/internal/service/model_factory.go b/internal/service/model_factory.go index 7ca8f53b10..b3b46edf19 100644 --- a/internal/service/model_factory.go +++ b/internal/service/model_factory.go @@ -29,6 +29,7 @@ import ( "ragflow/internal/dao" "ragflow/internal/entity" modelModule "ragflow/internal/entity/models" + "ragflow/internal/modelcontext" "ragflow/internal/permission" ) @@ -49,11 +50,31 @@ type modelTarget struct { SupportsTools bool } -// ModelAccess identifies the user and active tenant requesting a model. -// UserID may be empty for trusted background work scoped to TenantID. -type ModelAccess struct { - UserID string - TenantID string +// ErrNoDefaultModel reports that a membership has no configured model for the +// requested capability. +var ErrNoDefaultModel = errors.New("no default model configured") + +// ModelReference returns the internal ID when available, falling back to the +// composite name for built-in or catalog-only models. +func ModelReference(info *modelModule.ModelInfo) string { + if info == nil { + return "" + } + if id := strings.TrimSpace(info.ID); id != "" { + return id + } + name := strings.TrimSpace(info.Name) + if name == "" { + return "" + } + if provider := strings.TrimSpace(info.ProviderName); provider != "" { + instance := strings.TrimSpace(info.InstanceName) + if instance == "" { + return name + "@" + provider + } + return name + "@" + instance + "@" + provider + } + return name } // ModelFactory resolves model references, checks access, and constructs typed @@ -76,7 +97,7 @@ func NewModelFactory() *ModelFactory { } // NewChatModel resolves and constructs a chat model for the given access scope. -func (f *ModelFactory) NewChatModel(ctx context.Context, access ModelAccess, modelRef string) (*modelModule.ChatModel, error) { +func (f *ModelFactory) NewChatModel(ctx context.Context, access Identity, modelRef string) (*modelModule.ChatModel, error) { target, err := f.resolveConfig(ctx, access, entity.ModelTypeChat, modelRef) if err != nil { return nil, err @@ -85,7 +106,7 @@ func (f *ModelFactory) NewChatModel(ctx context.Context, access ModelAccess, mod } // NewDefaultChatModel constructs the tenant's configured default chat model. -func (f *ModelFactory) NewDefaultChatModel(ctx context.Context, access ModelAccess) (*modelModule.ChatModel, error) { +func (f *ModelFactory) NewDefaultChatModel(ctx context.Context, access Identity) (*modelModule.ChatModel, error) { target, err := f.resolveDefaultConfig(ctx, access, entity.ModelTypeChat) if err != nil { return nil, err @@ -94,7 +115,7 @@ func (f *ModelFactory) NewDefaultChatModel(ctx context.Context, access ModelAcce } // NewEmbeddingModel resolves and constructs an embedding model. -func (f *ModelFactory) NewEmbeddingModel(ctx context.Context, access ModelAccess, modelRef string) (*modelModule.EmbeddingModel, error) { +func (f *ModelFactory) NewEmbeddingModel(ctx context.Context, access Identity, modelRef string) (*modelModule.EmbeddingModel, error) { target, err := f.resolveConfig(ctx, access, entity.ModelTypeEmbedding, modelRef) if err != nil { return nil, err @@ -103,7 +124,7 @@ func (f *ModelFactory) NewEmbeddingModel(ctx context.Context, access ModelAccess } // NewDefaultEmbeddingModel constructs the tenant's configured default embedding model. -func (f *ModelFactory) NewDefaultEmbeddingModel(ctx context.Context, access ModelAccess) (*modelModule.EmbeddingModel, error) { +func (f *ModelFactory) NewDefaultEmbeddingModel(ctx context.Context, access Identity) (*modelModule.EmbeddingModel, error) { target, err := f.resolveDefaultConfig(ctx, access, entity.ModelTypeEmbedding) if err != nil { return nil, err @@ -112,7 +133,7 @@ func (f *ModelFactory) NewDefaultEmbeddingModel(ctx context.Context, access Mode } // NewRerankModel resolves and constructs a rerank model. -func (f *ModelFactory) NewRerankModel(ctx context.Context, access ModelAccess, modelRef string) (*modelModule.RerankModel, error) { +func (f *ModelFactory) NewRerankModel(ctx context.Context, access Identity, modelRef string) (*modelModule.RerankModel, error) { target, err := f.resolveConfig(ctx, access, entity.ModelTypeRerank, modelRef) if err != nil { return nil, err @@ -121,7 +142,7 @@ func (f *ModelFactory) NewRerankModel(ctx context.Context, access ModelAccess, m } // NewDefaultRerankModel constructs the tenant's configured default rerank model. -func (f *ModelFactory) NewDefaultRerankModel(ctx context.Context, access ModelAccess) (*modelModule.RerankModel, error) { +func (f *ModelFactory) NewDefaultRerankModel(ctx context.Context, access Identity) (*modelModule.RerankModel, error) { target, err := f.resolveDefaultConfig(ctx, access, entity.ModelTypeRerank) if err != nil { return nil, err @@ -130,7 +151,7 @@ func (f *ModelFactory) NewDefaultRerankModel(ctx context.Context, access ModelAc } // NewOCRModel resolves and constructs an OCR model. -func (f *ModelFactory) NewOCRModel(ctx context.Context, access ModelAccess, modelRef string) (*modelModule.OCRModel, error) { +func (f *ModelFactory) NewOCRModel(ctx context.Context, access Identity, modelRef string) (*modelModule.OCRModel, error) { target, err := f.resolveConfig(ctx, access, entity.ModelTypeOCR, modelRef) if err != nil { return nil, err @@ -139,7 +160,7 @@ func (f *ModelFactory) NewOCRModel(ctx context.Context, access ModelAccess, mode } // NewASRModel resolves and constructs a speech-to-text model. -func (f *ModelFactory) NewASRModel(ctx context.Context, access ModelAccess, modelRef string) (*modelModule.ASRModel, error) { +func (f *ModelFactory) NewASRModel(ctx context.Context, access Identity, modelRef string) (*modelModule.ASRModel, error) { target, err := f.resolveConfig(ctx, access, entity.ModelTypeSpeech2Text, modelRef) if err != nil { return nil, err @@ -148,7 +169,7 @@ func (f *ModelFactory) NewASRModel(ctx context.Context, access ModelAccess, mode } // NewDefaultASRModel constructs the tenant's configured default speech-to-text model. -func (f *ModelFactory) NewDefaultASRModel(ctx context.Context, access ModelAccess) (*modelModule.ASRModel, error) { +func (f *ModelFactory) NewDefaultASRModel(ctx context.Context, access Identity) (*modelModule.ASRModel, error) { target, err := f.resolveDefaultConfig(ctx, access, entity.ModelTypeSpeech2Text) if err != nil { return nil, err @@ -157,7 +178,7 @@ func (f *ModelFactory) NewDefaultASRModel(ctx context.Context, access ModelAcces } // NewTTSModel resolves and constructs a text-to-speech model. -func (f *ModelFactory) NewTTSModel(ctx context.Context, access ModelAccess, modelRef string) (*modelModule.TTSModel, error) { +func (f *ModelFactory) NewTTSModel(ctx context.Context, access Identity, modelRef string) (*modelModule.TTSModel, error) { target, err := f.resolveConfig(ctx, access, entity.ModelTypeTTS, modelRef) if err != nil { return nil, err @@ -166,7 +187,7 @@ func (f *ModelFactory) NewTTSModel(ctx context.Context, access ModelAccess, mode } // NewDefaultTTSModel constructs the tenant's configured default text-to-speech model. -func (f *ModelFactory) NewDefaultTTSModel(ctx context.Context, access ModelAccess) (*modelModule.TTSModel, error) { +func (f *ModelFactory) NewDefaultTTSModel(ctx context.Context, access Identity) (*modelModule.TTSModel, error) { target, err := f.resolveDefaultConfig(ctx, access, entity.ModelTypeTTS) if err != nil { return nil, err @@ -175,22 +196,54 @@ func (f *ModelFactory) NewDefaultTTSModel(ctx context.Context, access ModelAcces } // ResolveInfo returns model metadata without exposing API credentials or a driver. -func (f *ModelFactory) ResolveInfo(ctx context.Context, access ModelAccess, modelType entity.ModelType, modelRef string) (*modelModule.ModelInfo, error) { - target, err := f.resolveConfig(ctx, access, modelType, modelRef) +func (f *ModelFactory) ResolveInfo(ctx context.Context, access Identity, modelType entity.ModelType, modelRef string) (*modelModule.ModelInfo, error) { + access = identityFromContext(ctx, access) + if modelType == entity.ModelTypeOCR { + modelName, instanceName, providerName := splitRightAnchoredModelName(modelRef) + if modelName == "deepdoc" && instanceName == "default" && providerName == "infiniflow" { + if err := f.authorize(ctx, access); err != nil { + return nil, err + } + return &modelModule.ModelInfo{ + Name: modelName, + ProviderName: providerName, + InstanceName: instanceName, + ModelTypes: []string{entity.ModelTypeOCR.String()}, + }, nil + } + } + if err := f.authorize(ctx, access); err != nil { + return nil, err + } + modelRef = strings.TrimSpace(modelRef) + if name, ok := builtinEmbeddingName(modelRef); ok && modelType == entity.ModelTypeEmbedding { + info, _ := dao.GetModelProviderManager().GetModelByName("Builtin", name) + return modelInfoFromTarget(&modelTarget{ModelName: name, ModelType: modelType, ProviderName: "Builtin", ModelInfo: info}), nil + } + identity, err := f.resolveModelIdentity(ctx, access, modelType, modelRef) if err != nil { return nil, err } + target, err := modelTargetFromIdentity(identity, modelType) + if err != nil { + return nil, err + } + return modelInfoFromTarget(target), nil } // ResolveDefaultInfo returns metadata for the tenant's default model without // exposing its provider driver or API credentials. -func (f *ModelFactory) ResolveDefaultInfo(ctx context.Context, access ModelAccess, modelType entity.ModelType) (*modelModule.ModelInfo, error) { - target, err := f.resolveDefaultConfig(ctx, access, modelType) +func (f *ModelFactory) ResolveDefaultInfo(ctx context.Context, access Identity, modelType entity.ModelType) (*modelModule.ModelInfo, error) { + access = identityFromContext(ctx, access) + if err := f.authorize(ctx, access); err != nil { + return nil, err + } + ref, err := f.defaultModelReference(ctx, access, modelType) if err != nil { return nil, err } - return modelInfoFromTarget(target), nil + return f.ResolveInfo(ctx, access, modelType, ref) } func modelInfoFromTarget(target *modelTarget) *modelModule.ModelInfo { @@ -297,12 +350,12 @@ func splitRightAnchoredModelName(compositeName string) (modelName, instanceName, } type tenantModelExtra struct { - MaxTokens *int `json:"max_tokens"` - ModelTypes []string `json:"model_types"` - MaxDimension *int `json:"max_dimension"` - MaxBatchSize *int `json:"max_batch_size"` - Dimensions []int `json:"dimensions"` - Thinking *bool `json:"thinking"` + MaxTokens json.RawMessage `json:"max_tokens"` + ModelTypes []string `json:"model_types"` + MaxDimension *int `json:"max_dimension"` + MaxBatchSize *int `json:"max_batch_size"` + Dimensions []int `json:"dimensions"` + Thinking *bool `json:"thinking"` } func modelInfoWithTenantExtra(modelInfo *modelModule.Model, modelEntity *entity.TenantModel) (*modelModule.Model, error) { @@ -327,9 +380,13 @@ func modelInfoWithTenantExtra(modelInfo *modelModule.Model, modelEntity *entity. model.Thinking = &thinking } } + documentParseSupported := model.ModelTypeMap["doc_parse"] if modelEntity.ModelType != 0 { setModelInfoTypes(&model, entity.ModelType(modelEntity.ModelType)) + if documentParseSupported { + model.ModelTypeMap["doc_parse"] = true + } } if strings.TrimSpace(modelEntity.Extra) == "" { @@ -340,8 +397,8 @@ func modelInfoWithTenantExtra(modelInfo *modelModule.Model, modelEntity *entity. return nil, err } - if extra.MaxTokens != nil && *extra.MaxTokens > 0 { - contextLength := *extra.MaxTokens + if maxTokens := modelcontext.MaxTokensFromJSON(extra.MaxTokens); maxTokens > 0 { + contextLength := maxTokens model.ContextLength = &contextLength } if modelEntity.ModelType == 0 && len(extra.ModelTypes) > 0 { @@ -350,6 +407,9 @@ func modelInfoWithTenantExtra(modelInfo *modelModule.Model, modelEntity *entity. for _, modelType := range extra.ModelTypes { model.ModelTypeMap[modelType] = true } + if documentParseSupported { + model.ModelTypeMap["doc_parse"] = true + } } if extra.MaxDimension != nil && *extra.MaxDimension > 0 { model.MaxDimension = extra.MaxDimension @@ -404,33 +464,42 @@ func maxOutputFromModelInfo(modelInfo *modelModule.Model) int { return *modelInfo.MaxOutput } -func (f *ModelFactory) resolveConfig(ctx context.Context, access ModelAccess, modelType entity.ModelType, modelRef string) (*modelTarget, error) { - access = normalizeModelAccess(access) +func (f *ModelFactory) resolveConfig(ctx context.Context, access Identity, modelType entity.ModelType, modelRef string) (*modelTarget, error) { + access = identityFromContext(ctx, access) if err := f.authorize(ctx, access); err != nil { return nil, err } return f.resolveConfigAuthorized(ctx, access, modelType, modelRef) } -func (f *ModelFactory) resolveConfigAuthorized(ctx context.Context, access ModelAccess, modelType entity.ModelType, modelRef string) (*modelTarget, error) { +func (f *ModelFactory) resolveConfigAuthorized(ctx context.Context, access Identity, modelType entity.ModelType, modelRef string) (*modelTarget, error) { modelRef = strings.TrimSpace(modelRef) if modelRef == "" { return nil, fmt.Errorf("%w: model ref is required", errModelConfigUnavailable) } - if modelType == 0 { - return nil, fmt.Errorf("%w: model type is required", errModelConfigUnavailable) - } - + // A zero type skips category filtering for catalog-only capabilities such + // as doc_parse that are not represented by entity.ModelType. return f.resolveModel(ctx, access, modelType, modelRef) } -func normalizeModelAccess(access ModelAccess) ModelAccess { - access.TenantID = strings.TrimSpace(access.TenantID) - access.UserID = strings.TrimSpace(access.UserID) - return access +func normalizeIdentity(identity Identity) Identity { + identity.TenantID = strings.TrimSpace(identity.TenantID) + identity.UserID = strings.TrimSpace(identity.UserID) + return identity } -func (f *ModelFactory) authorize(ctx context.Context, access ModelAccess) error { +func identityFromContext(ctx context.Context, identity Identity) Identity { + identity = normalizeIdentity(identity) + if identity.UserID == "" { + if subject, ok := permission.SubjectFromContext(ctx); ok && strings.TrimSpace(subject.UserID) != "" && + (strings.TrimSpace(subject.TenantID) == "" || strings.TrimSpace(subject.TenantID) == identity.TenantID) { + identity.UserID = strings.TrimSpace(subject.UserID) + } + } + return identity +} + +func (f *ModelFactory) authorize(ctx context.Context, access Identity) error { if f == nil { return fmt.Errorf("%w: model factory is not initialized", errModelConfigUnavailable) } @@ -449,42 +518,54 @@ func (f *ModelFactory) authorize(ctx context.Context, access ModelAccess) error return nil } -func (f *ModelFactory) resolveDefaultConfig(ctx context.Context, access ModelAccess, modelType entity.ModelType) (*modelTarget, error) { - access = normalizeModelAccess(access) +func (f *ModelFactory) tenantOwnerMembership(ctx context.Context, tenantID string) (*entity.UserTenant, error) { + if f == nil || f.userTenantDAO == nil { + return nil, fmt.Errorf("%w: user tenant DAO is not initialized", errModelConfigUnavailable) + } + membership, err := f.userTenantDAO.GetOwnerByTenantID(ctx, dao.DB, tenantID) + if err != nil { + return nil, fmt.Errorf("failed to get tenant owner: %s: %w", tenantID, err) + } + return membership, nil +} + +func (f *ModelFactory) resolveDefaultConfig(ctx context.Context, access Identity, modelType entity.ModelType) (*modelTarget, error) { + access = identityFromContext(ctx, access) if err := f.authorize(ctx, access); err != nil { return nil, err } + ref, err := f.defaultModelReference(ctx, access, modelType) + if err != nil { + return nil, err + } + return f.resolveConfigAuthorized(ctx, access, modelType, ref) +} + +func (f *ModelFactory) defaultModelReference(ctx context.Context, access Identity, modelType entity.ModelType) (string, error) { if modelType == 0 { - return nil, fmt.Errorf("%w: model type is required", errModelConfigUnavailable) - } - if modelType == entity.ModelTypeOCR { - return nil, fmt.Errorf("OCR model name is required") - } - if f.userTenantDAO == nil { - return nil, fmt.Errorf("%w: user Tenant DAO is not initialized", errModelConfigUnavailable) + return "", fmt.Errorf("%w: model type is required", errModelConfigUnavailable) } var membership *entity.UserTenant var err error - if access.UserID != "" { + if access.UserID == "" { + membership, err = f.tenantOwnerMembership(ctx, access.TenantID) + } else { membership, err = f.userTenantDAO.FilterByUserIDAndTenantID(ctx, dao.DB, access.UserID, access.TenantID) - if err != nil && (access.UserID != access.TenantID || !errors.Is(err, gorm.ErrRecordNotFound)) { - return nil, fmt.Errorf("failed to get user tenant model settings: %s type %s: %w", access.TenantID, modelType, err) + if errors.Is(err, gorm.ErrRecordNotFound) && access.UserID == access.TenantID { + membership, err = f.tenantOwnerMembership(ctx, access.TenantID) } } - if membership == nil { - membership, err = f.userTenantDAO.GetOwnerByTenantID(ctx, dao.DB, access.TenantID) - if err != nil { - return nil, fmt.Errorf("failed to get tenant owner model settings: %s type %s: %w", access.TenantID, modelType, err) - } + if err != nil { + return "", fmt.Errorf("failed to get user tenant model settings: %s type %s: %w", access.TenantID, modelType, err) } modelName, modelID := defaultModelRefs(membership, modelType) if modelID != "" { - return f.resolveConfigAuthorized(ctx, access, modelType, modelID) + return modelID, nil } if strings.TrimSpace(modelName) == "" { - return nil, fmt.Errorf("no default %s model is set", modelType) + return "", fmt.Errorf("no default %s model is set: %w", modelType, ErrNoDefaultModel) } - return f.resolveConfigAuthorized(ctx, access, modelType, modelName) + return modelName, nil } func defaultModelRefs(membership *entity.UserTenant, modelType entity.ModelType) (string, string) { @@ -515,36 +596,24 @@ func defaultModelRefs(membership *entity.UserTenant, modelType entity.ModelType) // ResolveModelType returns every model category supported by modelRef. A // model may expose multiple categories, such as chat and vision. -func (s *ModelFactory) ResolveModelType(ctx context.Context, access ModelAccess, modelRef string) ([]entity.ModelType, error) { - access = normalizeModelAccess(access) +func (s *ModelFactory) ResolveModelType(ctx context.Context, access Identity, modelRef string) ([]entity.ModelType, error) { + access = identityFromContext(ctx, access) if err := s.authorize(ctx, access); err != nil { return nil, err } modelRef = strings.TrimSpace(modelRef) - if modelRef == "" { - return nil, fmt.Errorf("%w: model ref is required", errModelConfigUnavailable) + if _, ok := builtinEmbeddingName(modelRef); ok { + return []entity.ModelType{entity.ModelTypeEmbedding}, nil } - - model, err := s.lookupTenantModel(ctx, access.TenantID, modelRef) - if err == nil { - identity, identityErr := s.modelIdentity(ctx, access.TenantID, access.UserID, model, modelRef) - if identityErr != nil { - return nil, identityErr - } - return modelTypesFromBitmask(identity.modelType), nil - } - if !errors.Is(err, gorm.ErrRecordNotFound) { - return nil, err - } - - modelType, err := s.resolveCompositeModelType(ctx, access, modelRef) + identity, err := s.resolveModelIdentity(ctx, access, 0, modelRef) if err != nil { return nil, err } - return modelTypesFromBitmask(modelType), nil + return modelTypesFromBitmask(identity.modelType), nil } type modelIdentity struct { + modelName string modelEntity *entity.TenantModel providerEntity *entity.TenantModelProvider instanceEntity *entity.TenantModelInstance @@ -595,121 +664,154 @@ func catalogToolSupport(providerName, modelName string) bool { return model != nil && model.Tools != nil && model.Tools.Support } -func (s *ModelFactory) lookupTenantModel(ctx context.Context, tenantID, modelRef string) (*entity.TenantModel, error) { - if s == nil { +// resolveModelIdentity resolves saved IDs and composite names before applying the +// same ownership, relationship, status and resource permission checks. +func (s *ModelFactory) resolveModelIdentity(ctx context.Context, access Identity, requestedType entity.ModelType, modelRef string) (*modelIdentity, error) { + if s == nil || s.modelDAO == nil || s.modelProviderDAO == nil || s.modelInstanceDAO == nil { return nil, fmt.Errorf("%w: model factory is not initialized", errModelConfigUnavailable) } - if tenantID == "" { - return nil, fmt.Errorf("%w: tenant id is required", errModelConfigUnavailable) - } if modelRef == "" { return nil, fmt.Errorf("%w: model ref is required", errModelConfigUnavailable) } - if s.modelDAO == nil { - return nil, fmt.Errorf("%w: model factory is not initialized", errModelConfigUnavailable) - } - - return s.modelDAO.GetByID(ctx, dao.DB, modelRef) -} - -func (s *ModelFactory) modelIdentity(ctx context.Context, tenantID, userID string, modelEntity *entity.TenantModel, modelRef string) (*modelIdentity, error) { - if modelEntity == nil { - return nil, fmt.Errorf("%w: tenant model %q not found", errModelConfigUnavailable, modelRef) - } - - modelType := entity.ModelType(modelEntity.ModelType) - if modelType == 0 { - return nil, fmt.Errorf("%w: tenant model %q has no model type", errModelConfigUnavailable, modelRef) - } - if s.modelProviderDAO == nil || s.modelInstanceDAO == nil { - return nil, fmt.Errorf("%w: model factory is not initialized", errModelConfigUnavailable) - } - - providerEntity, err := s.modelProviderDAO.GetByID(ctx, dao.DB, modelEntity.ProviderID) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, fmt.Errorf("%w: provider id=%s not found for model %q", errModelConfigUnavailable, modelEntity.ProviderID, modelRef) - } + model, err := s.modelDAO.GetByID(ctx, dao.DB, modelRef) + if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { return nil, err } - if providerEntity == nil { - return nil, fmt.Errorf("%w: provider id=%s not found for model %q", errModelConfigUnavailable, modelEntity.ProviderID, modelRef) - } - if userID != "" { - if err := checkModelPermission(ctx, userID, permission.ResourceKindModelProvider, providerEntity.ID, permission.OperationRead); err != nil { - return nil, fmt.Errorf("%w: user %s cannot read model provider %q: %w", errModelConfigUnavailable, userID, providerEntity.ProviderName, err) - } - if err := checkModelPermission(ctx, userID, permission.ResourceKindModel, modelEntity.ID, permission.OperationUse); err != nil { - return nil, fmt.Errorf("%w: user %s cannot use model %q: %w", errModelConfigUnavailable, userID, modelRef, err) - } - } else if providerEntity.TenantID != tenantID { - return nil, fmt.Errorf( - "%w: tenant %s has no access to provider owned by tenant %s", - errModelConfigUnavailable, - tenantID, - providerEntity.TenantID, - ) - } - instanceEntity, err := s.modelInstanceDAO.GetByID(ctx, dao.DB, modelEntity.InstanceID) + owner, err := s.tenantOwnerMembership(ctx, access.TenantID) if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, fmt.Errorf("%w: instance id=%s not found for model %q", errModelConfigUnavailable, modelEntity.InstanceID, modelRef) - } return nil, err } - if instanceEntity == nil || instanceEntity.ProviderID != providerEntity.ID { - return nil, fmt.Errorf("%w: instance id=%s does not belong to provider id=%s", errModelConfigUnavailable, modelEntity.InstanceID, providerEntity.ID) - } - if userID != "" { - if err := checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instanceEntity.ID, permission.OperationUse); err != nil { - return nil, fmt.Errorf("%w: user %s cannot use model instance %q: %w", errModelConfigUnavailable, userID, instanceEntity.InstanceName, err) + identity := &modelIdentity{modelEntity: model} + var candidates []*entity.TenantModel + if model != nil { + identity.modelName = model.ModelName + identity.modelType = entity.ModelType(model.ModelType) + identity.providerEntity, err = s.modelProviderDAO.GetByID(ctx, dao.DB, model.ProviderID) + if err != nil { + return nil, fmt.Errorf("%w: provider id=%s lookup failed: %w", errModelConfigUnavailable, model.ProviderID, err) + } + identity.instanceEntity, err = s.modelInstanceDAO.GetByID(ctx, dao.DB, model.InstanceID) + if err != nil { + return nil, fmt.Errorf("%w: instance id=%s lookup failed: %w", errModelConfigUnavailable, model.InstanceID, err) + } + } else { + name, instanceName, providerName, parseErr := common.ParseCompositeModelName(modelRef) + if parseErr != nil { + return nil, fmt.Errorf("%w: model %q has no saved model row or valid composite name: %v", errModelConfigUnavailable, modelRef, parseErr) + } + identity.modelName = name + identity.providerEntity, err = s.modelProviderDAO.GetByTenantIDAndUserIDAndProviderName(ctx, dao.DB, access.TenantID, owner.UserID, providerName) + if err != nil { + return nil, fmt.Errorf("%w: provider %q lookup failed: %w", errModelConfigUnavailable, providerName, err) + } + identity.instanceEntity, err = s.modelInstanceDAO.GetByProviderIDAndInstanceName(ctx, dao.DB, identity.providerEntity.ID, instanceName) + if err != nil { + return nil, fmt.Errorf("%w: instance %q lookup failed: %w", errModelConfigUnavailable, instanceName, err) + } + candidates, err = s.modelDAO.GetModelsByProviderIDAndInstanceIDAndModelName(ctx, dao.DB, identity.providerEntity.ID, identity.instanceEntity.ID, name) + if err != nil { + return nil, fmt.Errorf("model %q lookup failed: %w", modelRef, err) + } + for _, candidate := range candidates { + if !strings.EqualFold(candidate.Status, "active") { + continue + } + identity.modelType |= entity.ModelType(candidate.ModelType) + if identity.modelEntity == nil && (requestedType == 0 || entity.ModelType(candidate.ModelType).Has(requestedType)) { + identity.modelEntity = candidate + } + } + if len(candidates) == 0 { + catalog, catalogErr := dao.GetModelProviderManager().GetModelByName(modelModule.CatalogProviderName(providerName), name) + if catalogErr == nil && catalog != nil { + identity.modelType = entity.ModelTypeFromStrings(catalog.ModelTypes) + } + } + if requestedType != 0 && identity.modelType != 0 && !identity.modelType.Has(requestedType) { + return nil, fmt.Errorf("%w: model %q cannot be used as %s model", errModelConfigUnavailable, modelRef, requestedType.String()) + } + if identity.modelEntity == nil && (len(candidates) > 0 || identity.modelType == 0) { + return nil, fmt.Errorf("%w: model %q is disabled or unavailable", errModelConfigUnavailable, modelRef) } } - if !strings.EqualFold(modelEntity.Status, "active") { + provider, instance := identity.providerEntity, identity.instanceEntity + if provider.TenantID != access.TenantID || provider.UserID != owner.UserID { + return nil, fmt.Errorf("%w: tenant %s can use only its owner's model configuration", errModelConfigUnavailable, access.TenantID) + } + if instance.ProviderID != provider.ID { + return nil, fmt.Errorf("%w: instance id=%s does not belong to provider id=%s", errModelConfigUnavailable, instance.ID, provider.ID) + } + if !strings.EqualFold(instance.Status, "active") { + return nil, fmt.Errorf("%w: model instance %q is disabled", errModelConfigUnavailable, instance.InstanceName) + } + if identity.modelEntity != nil && !strings.EqualFold(identity.modelEntity.Status, "active") { return nil, fmt.Errorf("%w: tenant model %q is disabled", errModelConfigUnavailable, modelRef) } - if !strings.EqualFold(instanceEntity.Status, "active") { - return nil, fmt.Errorf("%w: model instance %q is disabled", errModelConfigUnavailable, instanceEntity.InstanceName) - } - - return &modelIdentity{ - modelEntity: modelEntity, - providerEntity: providerEntity, - instanceEntity: instanceEntity, - modelType: modelType, - }, nil -} - -func (s *ModelFactory) resolveModel(ctx context.Context, access ModelAccess, modelType entity.ModelType, modelRef string) (*modelTarget, error) { - modelEntity, err := s.lookupTenantModel(ctx, access.TenantID, modelRef) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - // A bare ref that is BOTH unknown as a tenant model id AND rejected as - // a composite name is almost always a DANGLING REFERENCE — a knowledge - // base or chat pointing at a tenant_model row that was deleted. Say so - // explicitly: the underlying "provider name missing in model name: - // " reads like a naming-format mistake and sends the operator to - // look at the model's name instead of at the row that no longer exists. - composite, compositeErr := s.resolveProviderInstanceModel(ctx, access, modelType, modelRef) - if compositeErr != nil && !strings.Contains(modelRef, "@") { - return nil, fmt.Errorf("model %q is neither a tenant model id (no tenant_model row) nor a valid composite name — the reference is dangling: %w", modelRef, compositeErr) + if access.UserID != "" { + checker := permission.NewDatabaseChecker(dao.DB) + subject := permission.Subject{UserID: access.UserID} + if err := checker.CheckResource(ctx, subject, permission.ResourceRef{Kind: permission.ResourceKindModelProvider, ID: provider.ID}, permission.OperationRead); err != nil { + return nil, fmt.Errorf("%w: user %s cannot read model provider %q: %w", errModelConfigUnavailable, access.UserID, provider.ProviderName, err) + } + if err := checker.CheckResource(ctx, subject, permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, permission.OperationUse); err != nil { + return nil, fmt.Errorf("%w: user %s cannot use model instance %q: %w", errModelConfigUnavailable, access.UserID, instance.InstanceName, err) + } + checkedModels := candidates + if requestedType != 0 || model != nil { + checkedModels = nil + if identity.modelEntity != nil { + checkedModels = []*entity.TenantModel{identity.modelEntity} } - return composite, compositeErr + } + for _, candidate := range checkedModels { + if !strings.EqualFold(candidate.Status, "active") { + continue + } + if err := checker.CheckResource(ctx, subject, permission.ResourceRef{Kind: permission.ResourceKindModel, ID: candidate.ID}, permission.OperationUse); err != nil { + return nil, fmt.Errorf("%w: user %s cannot use model %q: %w", errModelConfigUnavailable, access.UserID, modelRef, err) + } } - return nil, err } - identity, err := s.modelIdentity(ctx, access.TenantID, access.UserID, modelEntity, modelRef) + if requestedType != 0 && !identity.modelType.Has(requestedType) { + return nil, fmt.Errorf("%w: model %q cannot be used as %s model", errModelConfigUnavailable, modelRef, requestedType.String()) + } + return identity, nil +} + +func (s *ModelFactory) resolveModel(ctx context.Context, access Identity, modelType entity.ModelType, modelRef string) (*modelTarget, error) { + if name, ok := builtinEmbeddingName(modelRef); ok && modelType == entity.ModelTypeEmbedding { + driver := modelModule.GetBuiltinEmbeddingModel(name) + if driver == nil { + return nil, fmt.Errorf("%w: builtin embedding model %q not found", errModelConfigUnavailable, name) + } + info, _ := dao.GetModelProviderManager().GetModelByName("Builtin", name) + config := &modelModule.APIConfig{} + if strings.Contains(common.GetEnv(common.EnvComposeProfiles), "tei-") && name == common.GetEnv(common.EnvTEIModel) { + baseURL := common.GetEnv(common.EnvTEIBaseURL) + config.BaseURL = &baseURL + } + return builtinModelTarget(modelType, name, driver, config, info), nil + } + identity, err := s.resolveModelIdentity(ctx, access, modelType, modelRef) if err != nil { return nil, err } - if !identity.modelType.Has(modelType) { - return nil, fmt.Errorf("%w: tenant model %q cannot be used as %s model", errModelConfigUnavailable, modelRef, modelType.String()) - } - return s.resolveTenantModel(ctx, modelType, modelRef, identity) + return s.resolveTenantModel(ctx, modelType, identity) } -func (s *ModelFactory) resolveTenantModel(ctx context.Context, modelType entity.ModelType, modelRef string, identity *modelIdentity) (*modelTarget, error) { - modelEntity := identity.modelEntity +func builtinEmbeddingName(modelRef string) (string, bool) { + if strings.Contains(common.GetEnv(common.EnvComposeProfiles), "tei-") { + name := common.GetEnv(common.EnvTEIModel) + pure, _, provider := splitRightAnchoredModelName(modelRef) + if name != "" && pure == name && (provider == "Builtin" || provider == "") { + return name, true + } + } + name, _, provider, err := common.ParseCompositeModelName(modelRef) + return name, err == nil && strings.EqualFold(provider, "Builtin") +} + +func (s *ModelFactory) resolveTenantModel(ctx context.Context, modelType entity.ModelType, identity *modelIdentity) (*modelTarget, error) { providerEntity := identity.providerEntity instanceEntity := identity.instanceEntity @@ -719,10 +821,9 @@ func (s *ModelFactory) resolveTenantModel(ctx context.Context, modelType entity. return nil, fmt.Errorf("%w: provider %q driver not found", errModelConfigUnavailable, providerEntity.ProviderName) } - modelInfo, _ := dao.GetModelProviderManager().GetModelByName(catalogName, modelEntity.ModelName) - modelInfo, err := modelInfoWithTenantExtra(modelInfo, modelEntity) + target, err := modelTargetFromIdentity(identity, modelType) if err != nil { - return nil, fmt.Errorf("%w: read model metadata: %v", errModelConfigUnavailable, err) + return nil, err } extra, err := decodeModelInstanceExtra(instanceEntity.Extra) @@ -737,211 +838,45 @@ func (s *ModelFactory) resolveTenantModel(ctx context.Context, modelType entity. return nil, fmt.Errorf("%w: create model driver: %v", errModelConfigUnavailable, err) } + target.Driver = driver + target.APIConfig = &modelModule.APIConfig{ApiKey: &apiKey, Region: ®ion, BaseURL: &baseURL} + return target, nil +} + +func modelTargetFromIdentity(identity *modelIdentity, modelType entity.ModelType) (*modelTarget, error) { + modelEntity, providerEntity, instanceEntity := identity.modelEntity, identity.providerEntity, identity.instanceEntity + catalogName := modelModule.CatalogProviderName(providerEntity.ProviderName) + modelInfo, _ := dao.GetModelProviderManager().GetModelByName(catalogName, identity.modelName) + modelInfo, err := modelInfoWithTenantExtra(modelInfo, modelEntity) + if err != nil { + return nil, fmt.Errorf("%w: read model metadata: %v", errModelConfigUnavailable, err) + } + + if modelEntity != nil && identity.modelType != 0 { + documentParseSupported := modelInfo.ModelTypeMap["doc_parse"] + setModelInfoTypes(modelInfo, identity.modelType) + if documentParseSupported { + modelInfo.ModelTypeMap["doc_parse"] = true + } + } + + modelExtra := "" + modelID := "" + if modelEntity != nil { + modelExtra = modelEntity.Extra + modelID = modelEntity.ID + } return &modelTarget{ - ModelID: modelEntity.ID, - ModelName: modelEntity.ModelName, + ModelID: modelID, + ModelName: identity.modelName, ModelType: modelType, ProviderName: providerEntity.ProviderName, InstanceID: instanceEntity.ID, InstanceName: instanceEntity.InstanceName, - Driver: driver, - APIConfig: &modelModule.APIConfig{ApiKey: &apiKey, Region: ®ion, BaseURL: &baseURL}, ModelInfo: modelInfo, ContextLength: contextLengthFromModelInfo(modelInfo), MaxOutput: maxOutputFromModelInfo(modelInfo), - SupportsTools: toolSupportFromEnrollment(modelEntity.Extra, instanceEntity.APIKey, providerEntity.ProviderName, modelEntity.ModelName), - }, nil -} - -func (s *ModelFactory) resolveProviderInstanceModel(ctx context.Context, access ModelAccess, modelType entity.ModelType, modelRef string) (*modelTarget, error) { - // TEI builtin embedding short-circuit. - if modelType == entity.ModelTypeEmbedding && strings.Contains(common.GetEnv(common.EnvComposeProfiles), "tei-") { - teiModel := common.GetEnv(common.EnvTEIModel) - teiBaseURL := common.GetEnv(common.EnvTEIBaseURL) - if modelRef == teiModel { - driver := modelModule.GetBuiltinEmbeddingModel(modelRef) - if driver == nil { - return nil, fmt.Errorf("%w: builtin (TEI) embedding model %q not found", errModelConfigUnavailable, modelRef) - } - modelInfo, _ := dao.GetModelProviderManager().GetModelByName("Builtin", modelRef) - return builtinModelTarget(modelType, modelRef, driver, &modelModule.APIConfig{BaseURL: &teiBaseURL}, modelInfo), nil - } - teiPure, _, teiProvider := splitRightAnchoredModelName(modelRef) - if teiPure == teiModel && (teiProvider == "Builtin" || teiProvider == "") { - driver := modelModule.GetBuiltinEmbeddingModel(teiPure) - if driver == nil { - return nil, fmt.Errorf("%w: builtin (TEI) embedding model %q not found", errModelConfigUnavailable, teiPure) - } - modelInfo, _ := dao.GetModelProviderManager().GetModelByName("Builtin", teiPure) - return builtinModelTarget(modelType, teiPure, driver, &modelModule.APIConfig{BaseURL: &teiBaseURL}, modelInfo), nil - } - } - if s == nil || s.modelProviderDAO == nil || s.modelInstanceDAO == nil || s.modelDAO == nil { - return nil, fmt.Errorf("%w: model factory is not initialized", errModelConfigUnavailable) - } - - pureModelName, instanceName, providerName, err := common.ParseCompositeModelName(modelRef) - if err != nil { - return nil, fmt.Errorf("%w: %v", errModelConfigUnavailable, err) - } - if modelType == entity.ModelTypeEmbedding { - var builtinName string - switch { - case strings.HasSuffix(modelRef, "@default@Builtin"): - builtinName = strings.TrimSuffix(modelRef, "@default@Builtin") - case strings.HasSuffix(modelRef, "@Builtin"): - builtinName = strings.TrimSuffix(modelRef, "@Builtin") - } - if builtinName != "" { - driver := modelModule.GetBuiltinEmbeddingModel(builtinName) - if driver == nil { - return nil, fmt.Errorf("%w: builtin embedding model %q not found", errModelConfigUnavailable, builtinName) - } - apiKey, region := "", "" - modelInfo, _ := dao.GetModelProviderManager().GetModelByName("Builtin", builtinName) - return builtinModelTarget(modelType, builtinName, driver, &modelModule.APIConfig{ApiKey: &apiKey, Region: ®ion}, modelInfo), nil - } - } - - provider, err := s.modelProviderDAO.GetByTenantIDAndProviderName(ctx, dao.DB, access.TenantID, providerName) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, fmt.Errorf("%w: provider %q lookup failed: %w", errModelConfigUnavailable, providerName, err) - } - return nil, fmt.Errorf("provider %q lookup failed: %w", providerName, err) - } - if provider == nil { - return nil, fmt.Errorf("%w: provider %q not found for model %q", errModelConfigUnavailable, providerName, modelRef) - } - if access.UserID != "" { - if err = checkModelPermission(ctx, access.UserID, permission.ResourceKindModelProvider, provider.ID, permission.OperationRead); err != nil { - return nil, fmt.Errorf("%w: user %s cannot read model provider %q: %w", errModelConfigUnavailable, access.UserID, providerName, err) - } - } - instance, err := s.modelInstanceDAO.GetByProviderIDAndInstanceName(ctx, dao.DB, provider.ID, instanceName) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, fmt.Errorf("%w: instance %q lookup failed: %w", errModelConfigUnavailable, instanceName, err) - } - return nil, fmt.Errorf("instance %q lookup failed: %w", instanceName, err) - } - if instance == nil { - return nil, fmt.Errorf("%w: instance %q not found for model %q", errModelConfigUnavailable, instanceName, modelRef) - } - if access.UserID != "" { - if err = checkModelPermission(ctx, access.UserID, permission.ResourceKindModelInstance, instance.ID, permission.OperationUse); err != nil { - return nil, fmt.Errorf("%w: user %s cannot use model instance %q: %w", errModelConfigUnavailable, access.UserID, instanceName, err) - } - } - if !strings.EqualFold(instance.Status, "active") { - return nil, fmt.Errorf("%w: model instance %q is disabled", errModelConfigUnavailable, instanceName) - } - extra, err := decodeModelInstanceExtra(instance.Extra) - if err != nil { - return nil, fmt.Errorf("%w: decode model instance configuration: %v", errModelConfigUnavailable, err) - } - region, baseURL, apiKey := extra.Region, extra.BaseURL, instance.APIKey - modelEntities, err := s.modelDAO.GetModelsByProviderIDAndInstanceIDAndModelName(ctx, dao.DB, provider.ID, instance.ID, pureModelName) - if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { - return nil, fmt.Errorf("model %q lookup failed: %w", modelRef, err) - } - var availableTypes entity.ModelType - var modelEntity *entity.TenantModel - var disabledForRequestedType bool - for _, candidate := range modelEntities { - if candidate == nil { - continue - } - candidateTypes := entity.ModelType(candidate.ModelType) - if strings.EqualFold(candidate.Status, "active") { - availableTypes |= candidateTypes - if candidateTypes.Has(modelType) && modelEntity == nil { - modelEntity = candidate - } - } else if candidateTypes.Has(modelType) { - disabledForRequestedType = true - } - } - if availableTypes != 0 && !availableTypes.Has(modelType) { - return nil, fmt.Errorf("%w: model %q cannot be used as %s model", errModelConfigUnavailable, modelRef, modelType.String()) - } - if modelEntity == nil && availableTypes == 0 && disabledForRequestedType { - return nil, fmt.Errorf("%w: model %q is disabled", errModelConfigUnavailable, modelRef) - } - - catalogName := modelModule.CatalogProviderName(providerName) - providerInfo := dao.GetModelProviderManager().FindProvider(catalogName) - if availableTypes != 0 { - if access.UserID != "" { - if err = checkModelPermission(ctx, access.UserID, permission.ResourceKindModel, modelEntity.ID, permission.OperationUse); err != nil { - return nil, fmt.Errorf("%w: user %s cannot use model %q: %w", errModelConfigUnavailable, access.UserID, modelRef, err) - } - } - if providerInfo == nil { - return nil, fmt.Errorf("%w: provider %q driver not found", errModelConfigUnavailable, providerName) - } - driver, driverErr := newModelDriverForBaseURL(providerInfo.ModelDriver, providerName, region, baseURL) - if driverErr != nil { - return nil, fmt.Errorf("%w: create model driver: %v", errModelConfigUnavailable, driverErr) - } - modelInfo, _ := dao.GetModelProviderManager().GetModelByName(catalogName, pureModelName) - modelInfo, err = modelInfoWithTenantExtra(modelInfo, modelEntity) - if err != nil { - return nil, fmt.Errorf("%w: read model metadata: %v", errModelConfigUnavailable, err) - } - setModelInfoTypes(modelInfo, availableTypes) - return &modelTarget{ - ModelID: modelEntity.ID, - ModelName: modelEntity.ModelName, - ModelType: modelType, - ProviderName: provider.ProviderName, - InstanceID: instance.ID, - InstanceName: instance.InstanceName, - Driver: driver, - APIConfig: &modelModule.APIConfig{ApiKey: &apiKey, Region: ®ion, BaseURL: &baseURL}, - ModelInfo: modelInfo, - ContextLength: contextLengthFromModelInfo(modelInfo), - MaxOutput: maxOutputFromModelInfo(modelInfo), - SupportsTools: toolSupportFromEnrollment(modelEntity.Extra, apiKey, provider.ProviderName, modelEntity.ModelName), - }, nil - } - if providerInfo == nil { - return nil, fmt.Errorf("%w: model provider config not found: %s", errModelConfigUnavailable, providerName) - } - targetProvider := dao.GetModelProviderManager().FindProvider(catalogName) - if targetProvider == nil { - return nil, fmt.Errorf("%w: model provider config not found: %s", errModelConfigUnavailable, providerName) - } - var modelInfo *modelModule.Model - for _, candidate := range targetProvider.Models { - if strings.EqualFold(candidate.Name, pureModelName) { - modelInfo = candidate - break - } - } - if modelInfo == nil { - return nil, fmt.Errorf("%w: model config not found: %s", errModelConfigUnavailable, modelRef) - } - availableTypes = entity.ModelTypeFromStrings(modelInfo.ModelTypes) - if !availableTypes.Has(modelType) { - return nil, fmt.Errorf("%w: model %q cannot be used as %s model", errModelConfigUnavailable, modelRef, modelType.String()) - } - driver, err := newModelDriverForBaseURL(targetProvider.ModelDriver, providerName, region, baseURL) - if err != nil { - return nil, fmt.Errorf("%w: create model driver: %v", errModelConfigUnavailable, err) - } - return &modelTarget{ - ModelName: modelInfo.Name, - ModelType: modelType, - ProviderName: provider.ProviderName, - InstanceID: instance.ID, - InstanceName: instance.InstanceName, - Driver: driver, - APIConfig: &modelModule.APIConfig{ApiKey: &apiKey, Region: ®ion, BaseURL: &baseURL}, - ModelInfo: modelInfo, - ContextLength: contextLengthFromModelInfo(modelInfo), - MaxOutput: maxOutputFromModelInfo(modelInfo), - SupportsTools: toolSupportFromEnrollment("", apiKey, provider.ProviderName, modelInfo.Name), + SupportsTools: toolSupportFromEnrollment(modelExtra, instanceEntity.APIKey, providerEntity.ProviderName, identity.modelName), }, nil } @@ -963,100 +898,6 @@ func builtinModelTarget(modelType entity.ModelType, modelName string, driver mod } } -func (s *ModelFactory) resolveCompositeModelType(ctx context.Context, access ModelAccess, modelRef string) (entity.ModelType, error) { - if strings.Contains(common.GetEnv(common.EnvComposeProfiles), "tei-") { - teiModel := common.GetEnv(common.EnvTEIModel) - if modelRef == teiModel { - if modelModule.GetBuiltinEmbeddingModel(teiModel) == nil { - return 0, fmt.Errorf("%w: builtin (TEI) embedding model %q not found", errModelConfigUnavailable, teiModel) - } - return entity.ModelTypeEmbedding, nil - } - } - pureModelName, instanceName, providerName, err := common.ParseCompositeModelName(modelRef) - if err != nil { - return 0, fmt.Errorf("%w: %v", errModelConfigUnavailable, err) - } - - if strings.EqualFold(providerName, "Builtin") { - if modelModule.GetBuiltinEmbeddingModel(pureModelName) == nil { - return 0, fmt.Errorf("%w: builtin model %q not found", errModelConfigUnavailable, pureModelName) - } - return entity.ModelTypeEmbedding, nil - } - if s == nil || s.modelProviderDAO == nil || s.modelInstanceDAO == nil || s.modelDAO == nil { - return 0, fmt.Errorf("%w: model factory is not initialized", errModelConfigUnavailable) - } - - provider, err := s.modelProviderDAO.GetByTenantIDAndProviderName(ctx, dao.DB, access.TenantID, providerName) - if err != nil { - return 0, fmt.Errorf("%w: provider %q lookup failed: %w", errModelConfigUnavailable, providerName, err) - } - if provider == nil { - return 0, fmt.Errorf("%w: provider %q not found for model %q", errModelConfigUnavailable, providerName, modelRef) - } - if access.UserID != "" { - if err = checkModelPermission(ctx, access.UserID, permission.ResourceKindModelProvider, provider.ID, permission.OperationRead); err != nil { - return 0, fmt.Errorf("%w: user %s cannot read model provider %q: %w", errModelConfigUnavailable, access.UserID, providerName, err) - } - } - instance, err := s.modelInstanceDAO.GetByProviderIDAndInstanceName(ctx, dao.DB, provider.ID, instanceName) - if err != nil { - return 0, fmt.Errorf("%w: instance %q lookup failed: %w", errModelConfigUnavailable, instanceName, err) - } - if instance == nil { - return 0, fmt.Errorf("%w: instance %q not found for model %q", errModelConfigUnavailable, instanceName, modelRef) - } - if !strings.EqualFold(instance.Status, "active") { - return 0, fmt.Errorf("%w: model instance %q is disabled", errModelConfigUnavailable, instanceName) - } - if access.UserID != "" { - if err = checkModelPermission(ctx, access.UserID, permission.ResourceKindModelInstance, instance.ID, permission.OperationUse); err != nil { - return 0, fmt.Errorf("%w: user %s cannot use model instance %q: %w", errModelConfigUnavailable, access.UserID, instanceName, err) - } - } - - models, err := s.modelDAO.GetModelsByProviderIDAndInstanceIDAndModelName(ctx, dao.DB, provider.ID, instance.ID, pureModelName) - if err != nil { - return 0, fmt.Errorf("%w: model %q lookup failed: %v", errModelConfigUnavailable, modelRef, err) - } - var activeTypes, inactiveTypes entity.ModelType - for _, model := range models { - if model == nil { - continue - } - if strings.EqualFold(model.Status, "active") { - if access.UserID != "" { - if err = checkModelPermission(ctx, access.UserID, permission.ResourceKindModel, model.ID, permission.OperationUse); err != nil { - return 0, fmt.Errorf("%w: user %s cannot use model %q: %w", errModelConfigUnavailable, access.UserID, modelRef, err) - } - } - activeTypes |= entity.ModelType(model.ModelType) - } else { - inactiveTypes |= entity.ModelType(model.ModelType) - } - } - if activeTypes != 0 { - return activeTypes, nil - } - - providerInfo := dao.GetModelProviderManager().FindProvider(modelModule.CatalogProviderName(providerName)) - if providerInfo == nil { - return 0, fmt.Errorf("%w: model provider config not found: %s", errModelConfigUnavailable, providerName) - } - for _, model := range providerInfo.Models { - if model != nil && strings.EqualFold(model.Name, pureModelName) { - catalogTypes := entity.ModelTypeFromStrings(model.ModelTypes) - availableTypes := catalogTypes &^ inactiveTypes - if availableTypes != 0 { - return availableTypes, nil - } - break - } - } - return 0, fmt.Errorf("%w: model %q not found for provider %q", errModelConfigUnavailable, pureModelName, providerName) -} - func modelTypesFromBitmask(modelType entity.ModelType) []entity.ModelType { modelTypes := make([]entity.ModelType, 0, 8) for _, candidate := range []entity.ModelType{ diff --git a/internal/service/model_factory_identity_test.go b/internal/service/model_factory_identity_test.go new file mode 100644 index 0000000000..70af103a15 --- /dev/null +++ b/internal/service/model_factory_identity_test.go @@ -0,0 +1,142 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package service + +import ( + "errors" + "reflect" + "testing" + + "ragflow/internal/entity" +) + +func TestModelFactoryIDAndNameShareValidation(t *testing.T) { + for _, test := range []struct { + name string + table string + updates map[string]interface{} + }{ + {name: "active"}, + {name: "disabled model", table: "tenant_model", updates: map[string]interface{}{"status": "inactive"}}, + {name: "disabled instance", table: "tenant_model_instance", updates: map[string]interface{}{"status": "inactive"}}, + {name: "foreign tenant", table: "tenant_model_provider", updates: map[string]interface{}{"tenant_id": "foreign"}}, + {name: "foreign owner", table: "tenant_model_provider", updates: map[string]interface{}{"user_id": "foreign"}}, + {name: "wrong provider", table: "tenant_model_instance", updates: map[string]interface{}{"provider_id": "foreign"}}, + } { + t.Run(test.name, func(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + useModelProviderServiceTestDB(t, db) + seedModelProviderServiceScope(t, db) + if test.table != "" { + if err := db.Table(test.table).Where("1 = 1").Updates(test.updates).Error; err != nil { + t.Fatal(err) + } + } + ctx := serviceTestContext(t, "user-1", "user-1") + factory := NewModelFactory() + access := Identity{UserID: "user-1", TenantID: "user-1"} + for _, ref := range []string{"model-1", "gpt-test@default@OpenAI"} { + info, infoErr := factory.ResolveInfo(ctx, access, entity.ModelTypeChat, ref) + types, typeErr := factory.ResolveModelType(ctx, access, ref) + _, runErr := factory.resolveConfig(ctx, access, entity.ModelTypeChat, ref) + for _, err := range []error{infoErr, typeErr, runErr} { + if test.table != "" { + if !errors.Is(err, errModelConfigUnavailable) { + t.Fatalf("%s: error = %v, want unavailable", ref, err) + } + } else if err != nil { + t.Fatalf("%s: %v", ref, err) + } + } + if test.table == "" { + if info.ID != "model-1" { + t.Fatalf("%s: ID = %q", ref, info.ID) + } + if !reflect.DeepEqual(types, []entity.ModelType{entity.ModelTypeChat}) { + t.Fatalf("%s: types = %v", ref, types) + } + } + } + }) + } +} + +func TestModelFactoryMetadataDoesNotBuildExecutionConfig(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + useModelProviderServiceTestDB(t, db) + seedModelProviderServiceScope(t, db) + if err := db.Model(&entity.TenantModelInstance{}).Where("id = ?", "instance-1").Update("extra", "invalid JSON").Error; err != nil { + t.Fatal(err) + } + if err := db.Model(&entity.UserTenant{}).Where("id = ?", "user-tenant-1").Updates(map[string]interface{}{"llm_id": "gpt-test@default@OpenAI", "tenant_llm_id": "model-1"}).Error; err != nil { + t.Fatal(err) + } + factory := NewModelFactory() + access := Identity{UserID: "user-1", TenantID: "user-1"} + ctx := serviceTestContext(t, "user-1", "user-1") + for _, ref := range []string{"model-1", "gpt-test@default@OpenAI"} { + if _, err := factory.ResolveInfo(ctx, access, entity.ModelTypeChat, ref); err != nil { + t.Fatalf("metadata %s: %v", ref, err) + } + if _, err := factory.ResolveModelType(ctx, access, ref); err != nil { + t.Fatalf("types %s: %v", ref, err) + } + if _, err := factory.resolveConfig(ctx, access, entity.ModelTypeChat, ref); err == nil { + t.Fatalf("execution %s accepted malformed configuration", ref) + } + } + if _, err := factory.ResolveDefaultInfo(ctx, access, entity.ModelTypeChat); err != nil { + t.Fatal(err) + } +} + +func TestModelFactoryCompositeSelectsRequestedType(t *testing.T) { + db := setupModelProviderServiceTestDB(t) + useModelProviderServiceTestDB(t, db) + seedModelProviderServiceScope(t, db) + if err := db.Create(&entity.TenantModel{ + ID: "model-embedding", ProviderID: "provider-1", InstanceID: "instance-1", ModelName: "gpt-test", ModelType: int(entity.ModelTypeEmbedding), Status: "active", + }).Error; err != nil { + t.Fatal(err) + } + factory := NewModelFactory() + access := Identity{UserID: "user-1", TenantID: "user-1"} + ctx := serviceTestContext(t, "user-1", "user-1") + ref := "gpt-test@default@OpenAI" + types, err := factory.ResolveModelType(ctx, access, ref) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(types, []entity.ModelType{entity.ModelTypeChat, entity.ModelTypeEmbedding}) { + t.Fatalf("types = %v", types) + } + info, err := factory.ResolveInfo(ctx, access, entity.ModelTypeEmbedding, ref) + if err != nil { + t.Fatal(err) + } + if info.ID != "model-embedding" { + t.Fatalf("selected model = %q", info.ID) + } + if err := db.Model(&entity.TenantModel{}).Where("id = ?", "model-embedding").Update("status", "inactive").Error; err != nil { + t.Fatal(err) + } + for _, modelRef := range []string{"model-embedding", ref} { + if _, err := factory.ResolveInfo(ctx, access, entity.ModelTypeEmbedding, modelRef); !errors.Is(err, errModelConfigUnavailable) { + t.Fatalf("%s: disabled model error = %v", modelRef, err) + } + } +} diff --git a/internal/service/model_permission.go b/internal/service/model_permission.go deleted file mode 100644 index ed7c76ce4b..0000000000 --- a/internal/service/model_permission.go +++ /dev/null @@ -1,53 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package service - -import ( - "context" - - "ragflow/internal/common" - "ragflow/internal/dao" - "ragflow/internal/permission" - permissionresponse "ragflow/internal/permission/response" -) - -func checkModelPermission(ctx context.Context, userID string, kind permission.ResourceKind, id string, operation permission.Operation) error { - return permission.NewDatabaseChecker(dao.DB).CheckResource( - ctx, - permission.Subject{UserID: userID}, - permission.ResourceRef{Kind: kind, ID: id}, - operation, - ) -} - -func accessibleModelResourceIDs(ctx context.Context, userID string, kind permission.ResourceKind, parent permission.ResourceRef, operation permission.Operation, entry *permission.ResourceRef, entryOperation permission.Operation) ([]string, error) { - query := permission.ScopeQuery{Kind: kind, Parent: parent, Operation: operation} - if entry != nil { - query.Entry = entry - query.EntryOp = entryOperation - } - scope, err := permission.NewDatabaseChecker(dao.DB).Scope(ctx, permission.Subject{UserID: userID}, query) - if err != nil { - return nil, err - } - if scope.Mode != permission.ScopeIDs { - return []string{}, nil - } - return scope.ResourceIDs, nil -} -func normalizeModelPermissionError(err error) (common.ErrorCode, error) { - return permissionresponse.NormalizeHidden(err) -} diff --git a/internal/service/model_service.go b/internal/service/model_service.go index 5c9ed09094..d995212472 100644 --- a/internal/service/model_service.go +++ b/internal/service/model_service.go @@ -72,6 +72,7 @@ func NewModelProviderService() *ModelProviderService { modelInstanceDAO: dao.NewTenantModelInstanceDAO(), modelDAO: dao.NewTenantModelDAO(), tenantDAO: dao.NewTenantDAO(), + userTenantDAO: dao.NewUserTenantDAO(), modelGroupDAO: dao.NewTenantModelGroupDAO(), modelGroupMapDAO: dao.NewTenantModelGroupMapDAO(), } @@ -82,6 +83,7 @@ type ModelProviderService struct { modelInstanceDAO *dao.TenantModelInstanceDAO modelDAO *dao.TenantModelDAO tenantDAO *dao.TenantDAO + userTenantDAO *dao.UserTenantDAO modelGroupDAO *dao.TenantModelGroupDAO modelGroupMapDAO *dao.TenantModelGroupMapDAO } @@ -111,16 +113,16 @@ type CheckConnectionRequest struct { func (m *ModelProviderService) AddModelProvider(ctx context.Context, providerName, userID string) (common.ErrorCode, error) { providerName = strings.TrimSpace(providerName) if err := checkFeatureAction(ctx, userID, "model_provider", permission.ActionEnable); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } tenantID, err := requireCurrentTenantID(ctx, userID) if err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } - existing, err := m.modelProviderDAO.GetByTenantIDAndProviderName(ctx, dao.DB, tenantID, providerName) + existing, err := m.modelProviderDAO.GetByTenantIDAndUserIDAndProviderName(ctx, dao.DB, tenantID, userID, providerName) if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { return common.CodeServerError, err } @@ -149,16 +151,18 @@ func (m *ModelProviderService) ListProvidersOfTenant(ctx context.Context, userID if errors.Is(err, permission.ErrMembershipNotFound) { return []map[string]interface{}{}, common.CodeSuccess, nil } - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, code, normalizedErr } - providerRows, err := m.modelProviderDAO.GetByTenantID(ctx, dao.DB, ownedTenantID) + providerRows, err := m.modelProviderDAO.ListByTenantIDAndUserID(ctx, dao.DB, ownedTenantID, userID) if err != nil { return nil, common.CodeServerError, err } providerRefs := make([]permission.ResourceRef, 0, len(providerRows)) + providerByID := make(map[string]*entity.TenantModelProvider, len(providerRows)) for _, provider := range providerRows { providerRefs = append(providerRefs, permission.ResourceRef{Kind: permission.ResourceKindModelProvider, ID: provider.ID}) + providerByID[provider.ID] = provider } accessibleProviders, err := permission.NewDatabaseChecker(dao.DB).FilterResources( ctx, @@ -167,17 +171,20 @@ func (m *ModelProviderService) ListProvidersOfTenant(ctx context.Context, userID permission.OperationRead, ) if err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, code, normalizedErr } - result := make([]map[string]interface{}, 0, len(accessibleProviders)) + type listedProvider struct { + id string + info map[string]interface{} + } + listed := make([]listedProvider, 0, len(accessibleProviders)) seenProviderNames := make(map[string]struct{}, len(accessibleProviders)) for _, providerRef := range accessibleProviders { - providerID := providerRef.ID - providerRow, err := m.modelProviderDAO.GetByID(ctx, dao.DB, providerID) - if err != nil { - return nil, common.CodeServerError, err + providerRow := providerByID[providerRef.ID] + if providerRow == nil { + continue } providerName := providerRow.ProviderName // Mirror Python's list_providers tenant branch: silently skip system-excluded @@ -206,27 +213,30 @@ func (m *ModelProviderService) ListProvidersOfTenant(ctx context.Context, userID // override it with providerName (the DB value), which may differ in // case from historical ToLower writes. - // Set has_instance flag. Mirrors Python's: - // provider_obj = TenantModelProviderService.get_by_tenant_id_and_provider_name(tenant_id, name) - // has_instance = bool(provider_obj and TenantModelInstanceService.get_all_by_provider_id(provider_obj.id)) - provider["has_instance"] = m.providerHasInstance(ctx, providerRow.ID) - - result = append(result, provider) + listed = append(listed, listedProvider{id: providerRow.ID, info: provider}) seenProviderNames[providerKey] = struct{}{} } + providerIDs := make([]string, 0, len(listed)) + for _, provider := range listed { + providerIDs = append(providerIDs, provider.id) + } + providersWithInstances, err := m.modelInstanceDAO.ListProviderIDsWithInstances(ctx, dao.DB, providerIDs) + if err != nil { + return nil, common.CodeServerError, err + } + hasInstance := make(map[string]bool, len(providerIDs)) + for _, providerID := range providersWithInstances { + hasInstance[providerID] = true + } + result := make([]map[string]interface{}, 0, len(listed)) + for _, provider := range listed { + provider.info["has_instance"] = hasInstance[provider.id] + result = append(result, provider.info) + } return result, common.CodeSuccess, nil } -// providerHasInstance checks whether a permission-scoped provider has any configured instances. -func (m *ModelProviderService) providerHasInstance(ctx context.Context, providerID string) bool { - instances, err := m.modelInstanceDAO.GetAllInstancesByProviderID(ctx, dao.DB, providerID) - if err != nil { - return false - } - return len(instances) > 0 -} - // isExcludedTenantProvider returns true for system-pool names that the Python // implementation (api/apps/services/provider_api_service.py:108) intentionally // filters out when listing a tenant's providers. @@ -242,7 +252,7 @@ func (m *ModelProviderService) DeleteModelProvider(ctx context.Context, userID, provider, err := m.resolveProviderInTenant(ctx, userID, providerName, permission.OperationDelete) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } return common.CodeNotFound, fmt.Errorf("provider %s not found", providerName) @@ -254,8 +264,13 @@ func (m *ModelProviderService) DeleteModelProvider(ctx context.Context, userID, return common.CodeServerError, err } for _, instance := range instances { - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationDelete); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationDelete, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } models, modelErr := m.modelDAO.GetModelsByInstanceID(ctx, dao.DB, instance.ID) @@ -263,8 +278,13 @@ func (m *ModelProviderService) DeleteModelProvider(ctx context.Context, userID, return common.CodeServerError, modelErr } for _, model := range models { - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, model.ID, permission.OperationDelete); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModel, ID: model.ID}, + permission.OperationDelete, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } } @@ -282,10 +302,13 @@ func (m *ModelProviderService) DeleteModelProvider(ctx context.Context, userID, } } - _, err = m.modelProviderDAO.DeleteByTenantIDAndProviderName(ctx, dao.DB, provider.TenantID, provider.ProviderName) + deleted, err := m.modelProviderDAO.DeleteByID(ctx, dao.DB, provider.ID) if err != nil { return common.CodeServerError, err } + if deleted != 1 { + return common.CodeNotFound, fmt.Errorf("provider %s not found", providerName) + } return common.CodeSuccess, nil } @@ -314,7 +337,7 @@ func (m *ModelProviderService) ResolveSavedProviderModelConfig(ctx context.Conte provider, err := m.resolveProviderInTenant(ctx, userID, providerName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - _, normalizedErr := normalizeModelPermissionError(err) + _, normalizedErr := permissionresponse.NormalizeHidden(err) return SavedProviderModelConfig{}, normalizedErr } return SavedProviderModelConfig{}, err @@ -335,11 +358,11 @@ func (m *ModelProviderService) ResolveSavedProviderModelConfig(ctx context.Conte if err == nil { err = permission.ErrPermissionDenied } - _, normalizedErr := normalizeModelPermissionError(err) + _, normalizedErr := permissionresponse.NormalizeHidden(err) return SavedProviderModelConfig{}, normalizedErr } if access.Source != permission.AccessSourceOwner && hasSavedProviderInstanceOverrides(options.APIKeyConfig, options.BaseURL, options.Region) { - _, normalizedErr := normalizeModelPermissionError(permission.ErrPermissionDenied) + _, normalizedErr := permissionresponse.NormalizeHidden(permission.ErrPermissionDenied) return SavedProviderModelConfig{}, normalizedErr } @@ -687,61 +710,49 @@ func mergeSavedProviderAPIKeyConfig(existing string, submitted json.RawMessage) return string(encoded), nil } -func (m *ModelProviderService) getProviderByIDOrName(ctx context.Context, userID, providerIDOrName string, operation permission.Operation) (*entity.TenantModelProvider, error) { +func (m *ModelProviderService) findProviderByIDOrNameInTenant(ctx context.Context, tenantID, userID, providerIDOrName string, includeTenantOwner bool) (*entity.TenantModelProvider, error) { providerIDOrName = strings.TrimSpace(providerIDOrName) if providerIDOrName == "" { return nil, gorm.ErrRecordNotFound } provider, err := m.modelProviderDAO.GetByID(ctx, dao.DB, providerIDOrName) if err == nil { - if err := checkModelPermission(ctx, userID, permission.ResourceKindModelProvider, provider.ID, operation); err != nil { - return nil, err + if provider.TenantID != tenantID { + return nil, gorm.ErrRecordNotFound + } + if provider.UserID != userID { + if !includeTenantOwner { + return nil, gorm.ErrRecordNotFound + } + owner, ownerErr := m.userTenantDAO.GetOwnerByTenantID(ctx, dao.DB, tenantID) + if ownerErr != nil { + return nil, ownerErr + } + if provider.UserID != owner.UserID { + return nil, gorm.ErrRecordNotFound + } } return provider, nil } if !errors.Is(err, gorm.ErrRecordNotFound) { return nil, err } - providerIDs, scopeErr := accessibleModelResourceIDs(ctx, userID, permission.ResourceKindModelProvider, permission.ResourceRef{}, operation, nil, "") - if scopeErr != nil { - return nil, scopeErr + provider, err = m.modelProviderDAO.GetByTenantIDAndUserIDAndProviderName(ctx, dao.DB, tenantID, userID, providerIDOrName) + if err == nil || !errors.Is(err, gorm.ErrRecordNotFound) || !includeTenantOwner { + return provider, err } - for _, providerID := range providerIDs { - candidate, candidateErr := m.modelProviderDAO.GetByID(ctx, dao.DB, providerID) - if candidateErr != nil { - return nil, candidateErr - } - if strings.EqualFold(candidate.ProviderName, providerIDOrName) { - return candidate, nil - } - } - return nil, gorm.ErrRecordNotFound -} - -func (m *ModelProviderService) findProviderByIDOrNameInTenant(ctx context.Context, tenantID, providerIDOrName string) (*entity.TenantModelProvider, error) { - providerIDOrName = strings.TrimSpace(providerIDOrName) - if providerIDOrName == "" { - return nil, gorm.ErrRecordNotFound - } - providers, err := m.modelProviderDAO.GetByTenantID(ctx, dao.DB, tenantID) + owner, err := m.userTenantDAO.GetOwnerByTenantID(ctx, dao.DB, tenantID) if err != nil { return nil, err } - for _, provider := range providers { - if provider.ID == providerIDOrName { - return provider, nil - } + if owner.UserID == userID { + return nil, gorm.ErrRecordNotFound } - for _, provider := range providers { - if strings.EqualFold(provider.ProviderName, providerIDOrName) { - return provider, nil - } - } - return nil, gorm.ErrRecordNotFound + return m.modelProviderDAO.GetByTenantIDAndUserIDAndProviderName(ctx, dao.DB, tenantID, owner.UserID, providerIDOrName) } -// resolveProviderInTenant resolves the provider addressed by a route inside the -// request's active tenant and authorizes operation on that row. +// resolveProviderInTenant resolves the caller's provider, or the active tenant +// owner's provider for read-only use, then authorizes the requested operation. // // Provider names are only unique per tenant (idx_tenant_provider_user_unique), // so a lookup across every tenant a user belongs to can select a same-named @@ -751,11 +762,16 @@ func (m *ModelProviderService) resolveProviderInTenant(ctx context.Context, user if err != nil { return nil, err } - provider, err := m.findProviderByIDOrNameInTenant(ctx, tenantID, providerIDOrName) + provider, err := m.findProviderByIDOrNameInTenant(ctx, tenantID, userID, providerIDOrName, operation == permission.OperationRead) if err != nil { return nil, err } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelProvider, provider.ID, operation); err != nil { + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelProvider, ID: provider.ID}, + operation, + ); err != nil { return nil, err } return provider, nil @@ -819,7 +835,7 @@ func (m *ModelProviderService) CreateProviderInstance(ctx context.Context, provi provider, err := m.resolveProviderInTenant(ctx, userID, providerIDOrName, permission.OperationCreate) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } return common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerIDOrName) @@ -848,35 +864,26 @@ func (m *ModelProviderService) CreateProviderInstance(ctx context.Context, provi if err := m.modelInstanceDAO.Create(ctx, tx, tenantModelInstance); err != nil { return err } - return permission.NewDatabaseChecker(tx).GrantResourceOwner(ctx, permission.Subject{UserID: userID, TenantID: tenantID}, permission.ResourceRef{ + if err := permission.NewDatabaseChecker(tx).GrantResourceOwner(ctx, permission.Subject{UserID: userID, TenantID: tenantID}, permission.ResourceRef{ Kind: permission.ResourceKindModelInstance, ID: tenantModelInstance.ID, - }) + }); err != nil { + return err + } + return m.addModelsToNewInstance(ctx, tx, provider, tenantModelInstance, modelInfo) }) if err != nil { return common.CodeServerError, fmt.Errorf("fail to create model instance: %w", err) } - // Add models to the instance. The instance row above is already committed - // (TenantModelInstanceDAO.Create owns its own transaction), so a failure - // here has to be undone explicitly or the tenant keeps a half-populated - // instance and can no longer reuse the instance name. - if err = m.addModelsToNewInstance(ctx, provider, tenantModelInstance, modelInfo); err != nil { - if rollbackErr := m.rollbackCreatedInstance(ctx, instanceID); rollbackErr != nil { - common.Logger.Error("failed to roll back model instance after model creation failure", - zap.String("instance_id", instanceID), zap.Error(rollbackErr)) - } - return common.CodeServerError, err - } - return common.CodeSuccess, nil } // addModelsToNewInstance creates the models of a freshly created provider // instance, either the requested ones or the provider's factory defaults. -func (m *ModelProviderService) addModelsToNewInstance(ctx context.Context, provider *entity.TenantModelProvider, instance *entity.TenantModelInstance, modelInfo []CreateInstanceModelInfo) error { +func (m *ModelProviderService) addModelsToNewInstance(ctx context.Context, db *gorm.DB, provider *entity.TenantModelProvider, instance *entity.TenantModelInstance, modelInfo []CreateInstanceModelInfo) error { if len(modelInfo) > 0 { for _, model := range modelInfo { - if err := m.addModelToInstance(ctx, provider, instance, model); err != nil { + if err := m.addModelToInstance(ctx, db, provider, instance, model); err != nil { return err } } @@ -901,7 +908,7 @@ func (m *ModelProviderService) addModelsToNewInstance(ctx context.Context, provi } else if llm.MaxTokens != nil { contextLength = *llm.MaxTokens } - if err := m.addModelToInstance(ctx, provider, instance, CreateInstanceModelInfo{ + if err := m.addModelToInstance(ctx, db, provider, instance, CreateInstanceModelInfo{ ModelName: llm.Name, ModelTypes: llm.ModelTypes, MaxTokens: contextLength, @@ -913,18 +920,6 @@ func (m *ModelProviderService) addModelsToNewInstance(ctx context.Context, provi return nil } -// rollbackCreatedInstance removes a provider instance together with any models -// already created under it, so a failed instance creation leaves nothing behind. -func (m *ModelProviderService) rollbackCreatedInstance(ctx context.Context, instanceID string) error { - return dao.DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - if _, err := m.modelDAO.DeleteByInstanceIDs(ctx, tx, []string{instanceID}); err != nil { - return err - } - _, err := m.modelInstanceDAO.DeleteByIDs(ctx, tx, []string{instanceID}) - return err - }) -} - // CreateNameOnlyProviderInstance creates a provider instance with only a name, // skipping API key validation and model creation. func (m *ModelProviderService) CreateNameOnlyProviderInstance(ctx context.Context, providerIDOrName, instanceName, userID string) (common.ErrorCode, error) { @@ -937,7 +932,7 @@ func (m *ModelProviderService) CreateNameOnlyProviderInstance(ctx context.Contex provider, err := m.resolveProviderInTenant(ctx, userID, providerIDOrName, permission.OperationCreate) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } return common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerIDOrName) @@ -968,9 +963,9 @@ func (m *ModelProviderService) CreateNameOnlyProviderInstance(ctx context.Contex } // addModelToInstance creates a single model under the given provider instance. -func (m *ModelProviderService) addModelToInstance(ctx context.Context, provider *entity.TenantModelProvider, instance *entity.TenantModelInstance, model CreateInstanceModelInfo) error { +func (m *ModelProviderService) addModelToInstance(ctx context.Context, db *gorm.DB, provider *entity.TenantModelProvider, instance *entity.TenantModelInstance, model CreateInstanceModelInfo) error { // Check for duplicate model. - _, err := m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(ctx, dao.DB, provider.ID, instance.ID, model.ModelName) + _, err := m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(ctx, db, provider.ID, instance.ID, model.ModelName) if err == nil { return fmt.Errorf("model '%s' already exists for provider '%s' and instance '%s'", model.ModelName, provider.ProviderName, instance.InstanceName) } @@ -984,7 +979,7 @@ func (m *ModelProviderService) addModelToInstance(ctx context.Context, provider combinedType |= entity.ModelTypeFromString(t) } - // Build extra fields. Mirrors Python's extra_fields in add_model_to_instance. + // Build the model configuration. extraFields := map[string]interface{}{ "max_tokens": model.MaxTokens, } @@ -993,7 +988,7 @@ func (m *ModelProviderService) addModelToInstance(ctx context.Context, provider extraFields[k] = v } } - if err = m.createModelRecord(ctx, dao.DB, provider.ID, instance.ID, model.ModelName, combinedType, "active", extraFields); err != nil { + if err = m.createModelRecord(ctx, db, provider.ID, instance.ID, model.ModelName, combinedType, "active", extraFields); err != nil { return fmt.Errorf("fail to create model '%s': %w", model.ModelName, err) } @@ -1027,7 +1022,7 @@ func (m *ModelProviderService) ListProviderInstances(ctx context.Context, provid return []map[string]interface{}{}, common.CodeSuccess, nil } if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, code, normalizedErr } if errors.Is(err, gorm.ErrRecordNotFound) { @@ -1035,20 +1030,19 @@ func (m *ModelProviderService) ListProviderInstances(ctx context.Context, provid } return nil, common.CodeServerError, err } - providerRef := permission.ResourceRef{Kind: permission.ResourceKindModelProvider, ID: provider.ID} - instances, err := m.modelInstanceDAO.GetAllInstancesByProviderID(ctx, dao.DB, provider.ID) if err != nil { return nil, common.CodeServerError, err } - instanceIDs, err := accessibleModelResourceIDs(ctx, userID, permission.ResourceKindModelInstance, providerRef, permission.OperationRead, &providerRef, permission.OperationRead) - if err != nil { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr + instanceRefs := make([]permission.ResourceRef, 0, len(instances)) + for _, instance := range instances { + instanceRefs = append(instanceRefs, permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}) } - accessibleInstances := make(map[string]struct{}, len(instanceIDs)) - for _, instanceID := range instanceIDs { - accessibleInstances[instanceID] = struct{}{} + checker := permission.NewDatabaseChecker(dao.DB) + instanceAccesses, err := checker.ResolveAccesses(ctx, permission.Subject{UserID: userID}, instanceRefs) + if err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) + return nil, code, normalizedErr } // Always emit a non-nil slice so the JSON encoder serializes [] rather @@ -1057,19 +1051,11 @@ func (m *ModelProviderService) ListProviderInstances(ctx context.Context, provid // web/src/pages/user-setting/setting-model/...) and would otherwise // crash on a freshly created tenant. result := make([]map[string]interface{}, 0, len(instances)) - checker := permission.NewDatabaseChecker(dao.DB) for _, instance := range instances { - if _, ok := accessibleInstances[instance.ID]; !ok { + access, ok := instanceAccesses[permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}] + if !ok || !access.Allows(permission.OperationRead) { continue } - access, accessErr := checker.ResolveAccess(ctx, permission.Subject{UserID: userID}, permission.ResourceRef{ - Kind: permission.ResourceKindModelInstance, - ID: instance.ID, - }) - if accessErr != nil { - code, normalizedErr := normalizeModelPermissionError(accessErr) - return nil, code, normalizedErr - } // Parse extra to extract region var extraFields map[string]string if instance.Extra != "" { @@ -1104,11 +1090,11 @@ func (m *ModelProviderService) ShowProviderInstance(ctx context.Context, provide provider, err := m.resolveProviderInTenant(ctx, userID, providerName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, code, normalizedErr } if errors.Is(err, gorm.ErrRecordNotFound) { - code, normalizedErr := normalizeModelPermissionError(permission.ErrResourceNotFound) + code, normalizedErr := permissionresponse.NormalizeHidden(permission.ErrResourceNotFound) return nil, code, normalizedErr } return nil, common.CodeServerError, err @@ -1129,7 +1115,7 @@ func (m *ModelProviderService) ShowProviderInstance(ctx context.Context, provide if err == nil { err = permission.ErrPermissionDenied } - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, code, normalizedErr } @@ -1165,29 +1151,43 @@ func (m *ModelProviderService) ShowProviderInstance(ctx context.Context, provide return result, common.CodeSuccess, nil } -func (m *ModelProviderService) ShowInstanceBalance(ctx context.Context, providerName, instanceName, userID string) (map[string]interface{}, common.ErrorCode, error) { +// resolveAuthorizedInstance checks the provider route and the requested instance operation. +func (m *ModelProviderService) resolveAuthorizedInstance(ctx context.Context, userID, providerName, instanceName string, operation permission.Operation) (*entity.TenantModelProvider, *entity.TenantModelInstance, common.ErrorCode, error) { provider, err := m.resolveProviderInTenant(ctx, userID, providerName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr + code, normalizedErr := permissionresponse.NormalizeHidden(err) + return nil, nil, code, normalizedErr } if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerName) + return nil, nil, common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerName) } - return nil, common.CodeServerError, err + return nil, nil, common.CodeServerError, err } instance, err := m.resolveInstance(ctx, provider, instanceName) if err != nil { if errors.Is(err, errInstanceProviderMismatch) || errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeNotFound, instanceLookupError(provider, instanceName, err) + return nil, nil, common.CodeNotFound, instanceLookupError(provider, instanceName, err) } - return nil, common.CodeServerError, err + return nil, nil, common.CodeServerError, err } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationUse); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + operation, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) + return nil, nil, code, normalizedErr + } + return provider, instance, common.CodeSuccess, nil +} + +func (m *ModelProviderService) ShowInstanceBalance(ctx context.Context, providerName, instanceName, userID string) (map[string]interface{}, common.ErrorCode, error) { + provider, instance, code, err := m.resolveAuthorizedInstance(ctx, userID, providerName, instanceName, permission.OperationUse) + if err != nil { + return nil, code, err } providerName = provider.ProviderName @@ -1252,8 +1252,13 @@ func (m *ModelProviderService) CheckConnection(ctx context.Context, providerName if !strings.EqualFold(provider.ProviderName, providerName) { return common.CodeNotFound, errors.New("model instance not found") } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelProvider, provider.ID, permission.OperationRead); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelProvider, ID: provider.ID}, + permission.OperationRead, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } instanceAccess, err = permission.NewDatabaseChecker(dao.DB).ResolveAccess(ctx, permission.Subject{UserID: userID}, permission.ResourceRef{ @@ -1264,7 +1269,7 @@ func (m *ModelProviderService) CheckConnection(ctx context.Context, providerName if err == nil { err = permission.ErrPermissionDenied } - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } savedInstance = instance @@ -1288,7 +1293,7 @@ func (m *ModelProviderService) CheckConnection(ctx context.Context, providerName return common.CodeBadRequest, errors.New("saved credentials cannot be combined with submitted credentials") } if instanceAccess.Source != permission.AccessSourceOwner && (region != "" || baseURL != "" || hasSavedProviderAPIKeyConfig(req.APIKeyConfig)) { - code, normalizedErr := normalizeModelPermissionError(permission.ErrPermissionDenied) + code, normalizedErr := permissionresponse.NormalizeHidden(permission.ErrPermissionDenied) return code, normalizedErr } apiKey = savedInstance.APIKey @@ -1349,7 +1354,7 @@ func (m *ModelProviderService) CheckConnection(ctx context.Context, providerName if instanceID != "" && len(modelVerifyResult) > 0 { if dbErr := m.updateModelVerifyResults(ctx, userID, providerName, instanceID, modelVerifyResult); dbErr != nil { if permissionresponse.IsPermissionError(dbErr) { - code, normalizedErr := normalizeModelPermissionError(dbErr) + code, normalizedErr := permissionresponse.NormalizeHidden(dbErr) return code, normalizedErr } common.Logger.Error("failed to persist model verify results", zap.Error(dbErr)) @@ -1372,7 +1377,12 @@ func (m *ModelProviderService) updateModelVerifyResults(ctx context.Context, use if err != nil { return fmt.Errorf("instance %s not found: %w", instanceID, err) } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationUpdate); err != nil { + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationUpdate, + ); err != nil { return err } provider, err := m.modelProviderDAO.GetByID(ctx, dao.DB, instance.ProviderID) @@ -1382,7 +1392,12 @@ func (m *ModelProviderService) updateModelVerifyResults(ctx context.Context, use if !strings.EqualFold(provider.ProviderName, providerName) { return permission.ErrResourceNotFound } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelProvider, provider.ID, permission.OperationRead); err != nil { + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelProvider, ID: provider.ID}, + permission.OperationRead, + ); err != nil { return err } @@ -1396,7 +1411,12 @@ func (m *ModelProviderService) updateModelVerifyResults(ctx context.Context, use } return fmt.Errorf("failed to look up %s for verify update: %w", modelName, err) } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, modelObj.ID, permission.OperationUpdate); err != nil { + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModel, ID: modelObj.ID}, + permission.OperationUpdate, + ); err != nil { return err } @@ -1757,41 +1777,18 @@ func minimalPNG() []byte { } } -func (m *ModelProviderService) ListTasks(ctx context.Context, providerName, instanceName, userID string) ([]modelModule.ListTaskStatus, common.ErrorCode, error) { - provider, err := m.resolveProviderInTenant(ctx, userID, providerName, permission.OperationRead) - if err != nil { - if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr - } - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerName) - } - return nil, common.CodeServerError, err - } - - instance, err := m.resolveInstance(ctx, provider, instanceName) - if err != nil { - if errors.Is(err, errInstanceProviderMismatch) || errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeNotFound, instanceLookupError(provider, instanceName, err) - } - return nil, common.CodeServerError, err - } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationUse); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr - } - providerName = provider.ProviderName +func taskDriverConfig(provider *entity.TenantModelProvider, instance *entity.TenantModelInstance) (modelModule.ModelDriver, *modelModule.APIConfig, error) { + providerName := provider.ProviderName providerInfo := dao.GetModelProviderManager().FindProvider(providerName) if providerInfo == nil { - return nil, common.CodeServerError, fmt.Errorf("provider %s not found", providerName) + return nil, nil, fmt.Errorf("provider %s not found", providerName) } var extra map[string]string - err = json.Unmarshal([]byte(instance.Extra), &extra) + err := json.Unmarshal([]byte(instance.Extra), &extra) if err != nil { - return nil, common.CodeServerError, err + return nil, nil, err } apiConfig := &modelModule.APIConfig{ @@ -1807,10 +1804,22 @@ func (m *ModelProviderService) ListTasks(ctx context.Context, providerName, inst if baseURL, ok := extra["base_url"]; ok && baseURL != "" { driver, err = newModelDriverForBaseURL(driver, providerName, region, baseURL) if err != nil { - return nil, common.CodeServerError, err + return nil, nil, err } } + return driver, apiConfig, nil +} + +func (m *ModelProviderService) ListTasks(ctx context.Context, providerName, instanceName, userID string) ([]modelModule.ListTaskStatus, common.ErrorCode, error) { + provider, instance, code, err := m.resolveAuthorizedInstance(ctx, userID, providerName, instanceName, permission.OperationUse) + if err != nil { + return nil, code, err + } + driver, apiConfig, err := taskDriverConfig(provider, instance) + if err != nil { + return nil, common.CodeServerError, err + } var listTaskResponse []modelModule.ListTaskStatus listTaskResponse, err = driver.ListTasks(ctx, apiConfig) if err != nil { @@ -1820,59 +1829,14 @@ func (m *ModelProviderService) ListTasks(ctx context.Context, providerName, inst } func (m *ModelProviderService) ShowTask(ctx context.Context, providerName, instanceName, taskID, userID string) (*modelModule.TaskResponse, common.ErrorCode, error) { - provider, err := m.resolveProviderInTenant(ctx, userID, providerName, permission.OperationRead) + provider, instance, code, err := m.resolveAuthorizedInstance(ctx, userID, providerName, instanceName, permission.OperationUse) if err != nil { - if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr - } - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerName) - } - return nil, common.CodeServerError, err + return nil, code, err } - - instance, err := m.resolveInstance(ctx, provider, instanceName) - if err != nil { - if errors.Is(err, errInstanceProviderMismatch) || errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeNotFound, instanceLookupError(provider, instanceName, err) - } - return nil, common.CodeServerError, err - } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationUse); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr - } - providerName = provider.ProviderName - - providerInfo := dao.GetModelProviderManager().FindProvider(providerName) - if providerInfo == nil { - return nil, common.CodeServerError, fmt.Errorf("provider %s not found", providerName) - } - - var extra map[string]string - err = json.Unmarshal([]byte(instance.Extra), &extra) + driver, apiConfig, err := taskDriverConfig(provider, instance) if err != nil { return nil, common.CodeServerError, err } - - apiConfig := &modelModule.APIConfig{ - ApiKey: nil, - Region: nil, - } - - region := extra["region"] - apiConfig.Region = ®ion - apiConfig.ApiKey = &instance.APIKey - - driver := providerInfo.ModelDriver - if baseURL, ok := extra["base_url"]; ok && baseURL != "" { - driver, err = newModelDriverForBaseURL(driver, providerName, region, baseURL) - if err != nil { - return nil, common.CodeServerError, err - } - } - var taskResponse *modelModule.TaskResponse taskResponse, err = driver.ShowTask(ctx, taskID, apiConfig) if err != nil { @@ -1881,31 +1845,10 @@ func (m *ModelProviderService) ShowTask(ctx context.Context, providerName, insta return taskResponse, common.CodeSuccess, nil } -// ListTenantAddedModels returns the list of models the tenant has "added" -// across all of their provider instances. It is the Go port of Python's -// models_api_service.list_tenant_added_models -// (api/apps/services/models_api_service.py:300) and is the response -// contract for GET /api/v1/models — the endpoint that -// web/src/hooks/use-llm-request.tsx → useFetchAllAddedModels consumes. -// -// Per the Python algorithm, for each (provider × instance) we cross-reference the factory catalog (internal/entity/models/model.go -// ProviderManager.Providers) with the per-tenant overrides in -// tenant_model: -// -// active_model_types = tenant_model rows with status='active' -// inactive_model_types = tenant_model rows with status='inactive' -// factory_model_types = provider.Models[i].ModelTypes -// model_types = (factory ∪ active) \ inactive -// -// The Go port never WRITES to tenant_model, so in practice every model -// from the factory catalog is treated as added unless explicitly -// disabled (which today can only happen via SQL — the Go port has no -// enable/disable endpoint path that mutates tenant_model). This is -// intentional: the previous Go contract mistakenly routed /api/v1/models -// to ListTenantDefaultModels (which only enumerates the 6-7 default -// tenant fields and returned `[]` for any tenant without defaults), -// breaking the front-end's "View Models" list entirely. -func (m *ModelProviderService) ListTenantAddedModels(ctx context.Context, userID, requestedTenantID, modelTypeFilter string) ([]map[string]interface{}, common.ErrorCode, error) { +// ListModelsForUserInTenant returns models available to userID in the requested +// tenant context, filtered by the user's resource permissions. Provider and +// instance configuration comes from the tenant owner. +func (m *ModelProviderService) ListModelsForUserInTenant(ctx context.Context, userID, requestedTenantID, modelTypeFilter string) ([]map[string]interface{}, common.ErrorCode, error) { tenant, code, err := m.resolveModelListTenant(ctx, userID, requestedTenantID) if err != nil { return nil, code, err @@ -1937,32 +1880,27 @@ func (m *ModelProviderService) ListTenantAddedModels(ctx context.Context, userID modelTypeFilterBin = entity.ModelTypeFromString(modelTypeFilter) } - providers, err := m.modelProviderDAO.GetByTenantID(ctx, dao.DB, tenantID) + checker := permission.NewDatabaseChecker(dao.DB) + tenantOwner, err := m.userTenantDAO.GetOwnerByTenantID(ctx, dao.DB, tenantID) + if err != nil { + return nil, common.CodeServerError, fmt.Errorf("get tenant owner for model listing: %w", err) + } + providerScope, err := checker.QueryScope(ctx, permission.Subject{UserID: userID}, permission.ScopeQuery{ + Kind: permission.ResourceKindModelProvider, + Operation: permission.OperationRead, + }) + if err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) + return nil, code, normalizedErr + } + providerQuery := providerScope.Apply(dao.DB.Model(&entity.TenantModelProvider{})) + providers, err := m.modelProviderDAO.ListByTenantIDAndUserID(ctx, providerQuery, tenantID, tenantOwner.UserID) if err != nil { return nil, common.CodeServerError, err } if len(providers) == 0 { return []map[string]interface{}{}, common.CodeSuccess, nil } - accessibleProviderIDs, err := accessibleModelResourceIDs(ctx, userID, permission.ResourceKindModelProvider, permission.ResourceRef{}, permission.OperationRead, nil, "") - if err != nil { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr - } - providerAllowed := make(map[string]struct{}, len(accessibleProviderIDs)) - for _, providerID := range accessibleProviderIDs { - providerAllowed[providerID] = struct{}{} - } - filteredProviders := providers[:0] - for _, provider := range providers { - if _, ok := providerAllowed[provider.ID]; ok { - filteredProviders = append(filteredProviders, provider) - } - } - providers = filteredProviders - if len(providers) == 0 { - return []map[string]interface{}{}, common.CodeSuccess, nil - } providerIDs := make([]string, 0, len(providers)) providerInfoByID := make(map[string]*entity.TenantModelProvider, len(providers)) @@ -1971,7 +1909,7 @@ func (m *ModelProviderService) ListTenantAddedModels(ctx context.Context, userID providerInfoByID[p.ID] = p } - instances, err := m.modelInstanceDAO.GetByProviderIDs(ctx, dao.DB, providerIDs) + instances, err := m.modelInstanceDAO.ListSummariesByProviderIDs(ctx, dao.DB, providerIDs) if err != nil { return nil, common.CodeServerError, err } @@ -1979,24 +1917,32 @@ func (m *ModelProviderService) ListTenantAddedModels(ctx context.Context, userID return []map[string]interface{}{}, common.CodeSuccess, nil } - instanceIDs := make([]string, 0, len(instances)) - instanceInfoByID := make(map[string]*entity.TenantModelInstance, len(instances)) + instanceInfoByID := make(map[string]*dao.TenantModelInstanceSummary, len(instances)) for _, inst := range instances { - instanceIDs = append(instanceIDs, inst.ID) instanceInfoByID[inst.ID] = inst } - // Fetch tenant model records and filter by model_type if needed. - modelRecords, err := m.modelDAO.GetModelsByProviderIDsAndInstanceIDs(ctx, dao.DB, providerIDs, instanceIDs) + // Apply model visibility in SQL so hidden model rows are never loaded into + // the service layer. Model access for a tenant is limited to its owner's + // provider configuration; personal model settings remain in the provider + // management endpoints. + modelScope, err := checker.QueryScope(ctx, permission.Subject{UserID: userID}, permission.ScopeQuery{ + Kind: permission.ResourceKindModel, + Operation: permission.OperationRead, + }) + if err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) + return nil, code, normalizedErr + } + modelRecords, err := m.modelDAO.ListByTenantIDAndModelType(ctx, dao.DB, tenantID, int(modelTypeFilterBin), modelScope.Apply) if err != nil { return nil, common.CodeServerError, err } - // Repair records that have model_type == 0. Before the ModelTypeFromString - // fix that added "asr" support, models whose factory catalog used the - // canonical name "asr" (instead of the legacy "speech2text") were stored - // with model_type = 0. Re-derive the correct bitmask from the factory - // catalog and persist it so the model shows up in typed queries. + // Derive model types for legacy zero-type rows in this response. Before the + // ModelTypeFromString fix added "asr" support, models using the canonical + // catalog name were stored with model_type = 0. Keep this compatibility + // repair read-only; schema/data repair does not belong in a list request. providerManager := dao.GetModelProviderManager() for _, rec := range modelRecords { if rec.ModelType != 0 { @@ -2007,7 +1953,7 @@ func (m *ModelProviderService) ListTenantAddedModels(ctx context.Context, userID continue } factoryModel, lookupErr := providerManager.GetModelByName(provInfo.ProviderName, rec.ModelName) - if lookupErr != nil || len(factoryModel.ModelTypes) == 0 { + if lookupErr != nil || factoryModel == nil || len(factoryModel.ModelTypes) == 0 { continue } combinedType := entity.ModelType(0) @@ -2016,9 +1962,6 @@ func (m *ModelProviderService) ListTenantAddedModels(ctx context.Context, userID } if combinedType != 0 { rec.ModelType = int(combinedType) - if permissionErr := checkModelPermission(ctx, userID, permission.ResourceKindModel, rec.ID, permission.OperationUpdate); permissionErr == nil { - _ = m.modelDAO.UpdateByID(ctx, dao.DB, rec.ID, map[string]interface{}{"model_type": int(combinedType)}) - } } } @@ -2032,20 +1975,6 @@ func (m *ModelProviderService) ListTenantAddedModels(ctx context.Context, userID } else { targetRecords = modelRecords } - modelRefs := make([]permission.ResourceRef, 0, len(targetRecords)) - for _, record := range targetRecords { - modelRefs = append(modelRefs, permission.ResourceRef{Kind: permission.ResourceKindModel, ID: record.ID}) - } - accessibleModels, err := permission.NewDatabaseChecker(dao.DB).FilterResources(ctx, permission.Subject{UserID: userID}, modelRefs, permission.OperationRead) - if err != nil { - code, normalizedErr := normalizeModelPermissionError(err) - return nil, code, normalizedErr - } - accessibleModelIDs := make(map[string]struct{}, len(accessibleModels)) - for _, ref := range accessibleModels { - accessibleModelIDs[ref.ID] = struct{}{} - } - // Build model rank map from factory catalog (mirrors Python's model_rank_map). modelRankMap := make(map[string]int) // key: "providerName@modelName" factoryRankMapping := make(map[string]int) @@ -2063,9 +1992,6 @@ func (m *ModelProviderService) ListTenantAddedModels(ctx context.Context, userID added := make([]map[string]interface{}, 0, len(targetRecords)) for _, modelRecord := range targetRecords { - if _, ok := accessibleModelIDs[modelRecord.ID]; !ok { - continue - } provInfo, provOK := providerInfoByID[modelRecord.ProviderID] instInfo, instOK := instanceInfoByID[modelRecord.InstanceID] if !provOK || !instOK { @@ -2165,12 +2091,12 @@ func (m *ModelProviderService) resolveModelListTenant(ctx context.Context, userI if errors.Is(err, permission.ErrMembershipNotFound) { return nil, common.CodeSuccess, nil } - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, code, normalizedErr } } else { if err := permission.NewDatabaseChecker(dao.DB).CheckTenant(ctx, permission.Subject{UserID: userID}, tenantID, permission.TenantMember); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, code, normalizedErr } } @@ -2311,7 +2237,7 @@ func (m *ModelProviderService) ensureOCRProviderFromEnv(ctx context.Context, use } // 1. Find or create the provider. - provider, err := m.modelProviderDAO.GetByTenantIDAndProviderName(ctx, dao.DB, tenantID, providerName) + provider, err := m.modelProviderDAO.GetByTenantIDAndUserIDAndProviderName(ctx, dao.DB, tenantID, userID, providerName) if err != nil { if !dao.IsNotFoundErr(err) { return fmt.Errorf("failed to get provider %s: %w", providerName, err) @@ -2397,7 +2323,7 @@ func (m *ModelProviderService) AlterProviderInstance(ctx context.Context, userID provider, err := m.resolveProviderInTenant(ctx, userID, providerIDOrName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } return common.CodeNotFound, fmt.Errorf("provider '%s' does not exist", providerIDOrName) @@ -2411,8 +2337,13 @@ func (m *ModelProviderService) AlterProviderInstance(ctx context.Context, userID } return common.CodeNotFound, fmt.Errorf("no instance found for provider '%s' and instance '%s'", providerName, instanceIDOrName) } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationUpdate); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationUpdate, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } @@ -2438,8 +2369,13 @@ func (m *ModelProviderService) AlterProviderInstance(ctx context.Context, userID if !submittedModelNames[name] { operation = permission.OperationDelete } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, model.ID, operation); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModel, ID: model.ID}, + operation, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } } @@ -2448,8 +2384,13 @@ func (m *ModelProviderService) AlterProviderInstance(ctx context.Context, userID continue } if _, exists := existingModelMap[model.ModelName]; !exists { - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationCreate); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationCreate, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } break @@ -2552,7 +2493,7 @@ func (m *ModelProviderService) AlterProviderInstance(ctx context.Context, userID } } else { // Add new model. - if err = m.addModelToInstance(ctx, provider, instance, mdl); err != nil { + if err = m.addModelToInstance(ctx, dao.DB, provider, instance, mdl); err != nil { return common.CodeServerError, err } } @@ -2571,7 +2512,7 @@ func (m *ModelProviderService) DropProviderInstances(ctx context.Context, provid provider, err := m.resolveProviderInTenant(ctx, userID, providerIDOrName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } if errors.Is(err, gorm.ErrRecordNotFound) { @@ -2598,8 +2539,13 @@ func (m *ModelProviderService) DropProviderInstances(ctx context.Context, provid return common.CodeServerError, resolveErr } instanceIDs = append(instanceIDs, instance.ID) - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationDelete); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationDelete, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } models, modelErr := m.modelDAO.GetModelsByInstanceID(ctx, dao.DB, instance.ID) @@ -2607,8 +2553,13 @@ func (m *ModelProviderService) DropProviderInstances(ctx context.Context, provid return common.CodeServerError, modelErr } for _, model := range models { - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, model.ID, permission.OperationDelete); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModel, ID: model.ID}, + permission.OperationDelete, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } } @@ -2644,7 +2595,7 @@ func (m *ModelProviderService) DropInstanceModels(ctx context.Context, providerI provider, err := m.resolveProviderInTenant(ctx, userID, providerIDOrName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } return common.CodeDataError, fmt.Errorf("no provider found for provider '%s'", providerIDOrName) @@ -2657,8 +2608,13 @@ func (m *ModelProviderService) DropInstanceModels(ctx context.Context, providerI } return common.CodeDataError, fmt.Errorf("no instance found for provider '%s' and instance '%s'", providerIDOrName, instanceIDOrName) } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationRead); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationRead, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } @@ -2691,8 +2647,13 @@ func (m *ModelProviderService) DropInstanceModels(ctx context.Context, providerI idsToDelete := make([]string, 0, len(modelNames)) for _, obj := range modelObjs { if _, ok := requestedNames[obj.ModelName]; ok { - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, obj.ID, permission.OperationDelete); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModel, ID: obj.ID}, + permission.OperationDelete, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } idsToDelete = append(idsToDelete, obj.ID) @@ -2710,7 +2671,7 @@ func (m *ModelProviderService) ListInstanceModels(ctx context.Context, providerN provider, err := m.resolveProviderInTenant(ctx, userID, providerName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - _, normalizedErr := normalizeModelPermissionError(err) + _, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, normalizedErr } return nil, err @@ -2723,8 +2684,13 @@ func (m *ModelProviderService) ListInstanceModels(ctx context.Context, providerN } return nil, err } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationRead); err != nil { - _, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationRead, + ); err != nil { + _, normalizedErr := permissionresponse.NormalizeHidden(err) return nil, normalizedErr } @@ -2733,15 +2699,21 @@ func (m *ModelProviderService) ListInstanceModels(ctx context.Context, providerN if err != nil { return nil, err } + modelRefs := make([]permission.ResourceRef, 0, len(modelObjs)) + for _, model := range modelObjs { + modelRefs = append(modelRefs, permission.ResourceRef{Kind: permission.ResourceKindModel, ID: model.ID}) + } + resourceAccesses, err := permission.NewDatabaseChecker(dao.DB).ResolveAccesses(ctx, permission.Subject{UserID: userID}, modelRefs) + if err != nil { + _, normalizedErr := permissionresponse.NormalizeHidden(err) + return nil, normalizedErr + } modelList := make([]map[string]interface{}, 0, len(modelObjs)) for _, model := range modelObjs { - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, model.ID, permission.OperationRead); err != nil { - if errors.Is(err, permission.ErrPermissionDenied) || errors.Is(err, permission.ErrResourceNotFound) { - continue - } - _, normalizedErr := normalizeModelPermissionError(err) - return nil, normalizedErr + access, ok := resourceAccesses[permission.ResourceRef{Kind: permission.ResourceKindModel, ID: model.ID}] + if !ok || !access.Allows(permission.OperationRead) { + continue } modelExtra := make(map[string]interface{}) if model.Extra != "" { @@ -2809,7 +2781,7 @@ func (m *ModelProviderService) AlterModel(ctx context.Context, providerName, ins provider, err := m.resolveProviderInTenant(ctx, userID, providerName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } if errors.Is(err, gorm.ErrRecordNotFound) { @@ -2825,8 +2797,13 @@ func (m *ModelProviderService) AlterModel(ctx context.Context, providerName, ins } return common.CodeServerError, err } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationRead); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationRead, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } @@ -2854,8 +2831,13 @@ func (m *ModelProviderService) AlterModel(ctx context.Context, providerName, ins return common.CodeNotFound, fmt.Errorf("model with ID %q does not belong to provider %q and instance %q", modelID, providerName, instanceName) } } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, model.ID, permission.OperationUpdate); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModel, ID: model.ID}, + permission.OperationUpdate, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } @@ -2941,174 +2923,6 @@ func resolveModelType(val interface{}) int { } } -type ModelInstanceAndProviderInfo struct { - ProviderEntity *entity.TenantModelProvider - ProviderInfo *modelModule.Provider - InstanceEntity *entity.TenantModelInstance - ModelEntity *entity.TenantModel - ModelInfo *modelModule.Model - APIConfig *modelModule.APIConfig -} - -func (m *ModelProviderService) getModelInstanceAndProviderByName(ctx context.Context, providerName, instanceName, modelName *string, userID string, apiConfig *modelModule.APIConfig) (*ModelInstanceAndProviderInfo, error) { - if providerName == nil || instanceName == nil || modelName == nil { - return nil, errors.New("provider name, instance name and model name are required when model id is absent") - } - - // Provider and instance names are only unique inside one tenant, so a - // composite reference resolves inside the request's active tenant. - tenantID, err := requireCurrentTenantID(ctx, userID) - if err != nil { - return nil, err - } - providerEntity, err := m.findProviderByIDOrNameInTenant(ctx, tenantID, *providerName) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, permission.ErrResourceNotFound - } - return nil, err - } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelProvider, providerEntity.ID, permission.OperationRead); err != nil { - return nil, err - } - - instanceEntity, err := m.resolveInstance(ctx, providerEntity, *instanceName) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) || errors.Is(err, errInstanceProviderMismatch) { - return nil, permission.ErrResourceNotFound - } - return nil, err - } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instanceEntity.ID, permission.OperationUse); err != nil { - return nil, err - } - - modelEntity, err := m.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(ctx, dao.DB, providerEntity.ID, instanceEntity.ID, *modelName) - if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { - return nil, err - } - if errors.Is(err, gorm.ErrRecordNotFound) { - modelEntity = nil - } - if modelEntity != nil { - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, modelEntity.ID, permission.OperationUse); err != nil { - return nil, err - } - } - - catalogName := modelModule.CatalogProviderName(providerEntity.ProviderName) - providerInfo := dao.GetModelProviderManager().FindProvider(catalogName) - if providerInfo == nil { - return nil, errors.New("provider not found") - } - - modelInfo, err := dao.GetModelProviderManager().GetModelByName(catalogName, *modelName) - if err != nil { - return nil, fmt.Errorf("provider %s model %s not found", providerEntity.ProviderName, *modelName) - } - modelInfo, err = modelInfoWithTenantExtra(modelInfo, modelEntity) - if err != nil { - return nil, err - } - - var extra map[string]string - err = json.Unmarshal([]byte(instanceEntity.Extra), &extra) - if err != nil { - return nil, err - } - - if apiConfig == nil { - apiConfig = &modelModule.APIConfig{} - } - - region := extra["region"] - baseURL := extra["base_url"] - - apiConfig.ApiKey = &instanceEntity.APIKey - apiConfig.BaseURL = &baseURL - apiConfig.Region = ®ion - - var result = &ModelInstanceAndProviderInfo{ - ProviderEntity: providerEntity, - ProviderInfo: providerInfo, - InstanceEntity: instanceEntity, - ModelEntity: modelEntity, - ModelInfo: modelInfo, - APIConfig: apiConfig, - } - - return result, nil -} - -func (m *ModelProviderService) getModelInstanceAndProviderByID(ctx context.Context, modelID *string, userID string, apiConfig *modelModule.APIConfig) (*ModelInstanceAndProviderInfo, error) { - modelEntity, err := m.modelDAO.GetByID(ctx, dao.DB, *modelID) - if err != nil { - return nil, err - } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModel, modelEntity.ID, permission.OperationUse); err != nil { - return nil, err - } - - instanceEntity, err := m.modelInstanceDAO.GetByID(ctx, dao.DB, modelEntity.InstanceID) - if err != nil { - return nil, err - } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instanceEntity.ID, permission.OperationUse); err != nil { - return nil, err - } - - providerEntity, err := m.modelProviderDAO.GetByID(ctx, dao.DB, instanceEntity.ProviderID) - if err != nil { - return nil, err - } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelProvider, providerEntity.ID, permission.OperationRead); err != nil { - return nil, err - } - - catalogName := modelModule.CatalogProviderName(providerEntity.ProviderName) - providerInfo := dao.GetModelProviderManager().FindProvider(catalogName) - if providerInfo == nil { - return nil, errors.New("provider not found") - } - - modelInfo, err := dao.GetModelProviderManager().GetModelByName(catalogName, modelEntity.ModelName) - if err != nil { - return nil, errors.New(fmt.Sprintf("provider %s model %s not found", providerEntity.ProviderName, modelEntity.ModelName)) - } - modelInfo, err = modelInfoWithTenantExtra(modelInfo, modelEntity) - if err != nil { - return nil, err - } - - var extra map[string]string - err = json.Unmarshal([]byte(instanceEntity.Extra), &extra) - if err != nil { - return nil, err - } - - if apiConfig == nil { - apiConfig = &modelModule.APIConfig{} - } - - region := extra["region"] - baseURL := extra["base_url"] - - apiConfig.ApiKey = &instanceEntity.APIKey - apiConfig.BaseURL = &baseURL - apiConfig.Region = ®ion - - var result = &ModelInstanceAndProviderInfo{ - ProviderEntity: providerEntity, - ProviderInfo: providerInfo, - InstanceEntity: instanceEntity, - ModelEntity: modelEntity, - ModelInfo: modelInfo, - APIConfig: apiConfig, - } - - return result, nil -} - func validateEmbeddingModel(model *modelModule.Model, requestedDimension, requestedBatchSize int) error { if model == nil { return fmt.Errorf("embedding model is nil") @@ -3204,7 +3018,7 @@ func (m *ModelProviderService) AddModel(ctx context.Context, request *AddModelRe provider, err := m.resolveProviderInTenant(ctx, userID, request.ProviderName, permission.OperationRead) if err != nil { if permissionresponse.IsPermissionError(err) { - code, normalizedErr := normalizeModelPermissionError(err) + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } return common.CodeDataError, fmt.Errorf("no provider found for provider '%s'", request.ProviderName) @@ -3216,8 +3030,13 @@ func (m *ModelProviderService) AddModel(ctx context.Context, request *AddModelRe } return common.CodeDataError, fmt.Errorf("no instance found for provider '%s' and instance '%s'", request.ProviderName, request.InstanceName) } - if err = checkModelPermission(ctx, userID, permission.ResourceKindModelInstance, instance.ID, permission.OperationCreate); err != nil { - code, normalizedErr := normalizeModelPermissionError(err) + if err = permission.NewDatabaseChecker(dao.DB).CheckResource( + ctx, + permission.Subject{UserID: userID}, + permission.ResourceRef{Kind: permission.ResourceKindModelInstance, ID: instance.ID}, + permission.OperationCreate, + ); err != nil { + code, normalizedErr := permissionresponse.NormalizeHidden(err) return code, normalizedErr } diff --git a/internal/service/model_service_test.go b/internal/service/model_service_test.go index 037518e4c4..923d82c4f1 100644 --- a/internal/service/model_service_test.go +++ b/internal/service/model_service_test.go @@ -574,8 +574,13 @@ func TestCreateProviderInstanceRollsBackInstanceWhenModelCreationFails(t *testin useModelProviderServiceTestDB(t, db) seedModelProviderServiceScope(t, db) - // The second entry duplicates the first one, so model creation fails after - // the instance row has already been committed. + // Reject compensating DELETEs: failure must roll back the transaction, + // including the first model, without issuing any cleanup statements. + if err := db.Exec(`CREATE TRIGGER reject_instance_cleanup BEFORE DELETE ON tenant_model_instance + BEGIN SELECT RAISE(ABORT, 'instance cleanup is forbidden'); END`).Error; err != nil { + t.Fatal(err) + } + // The second entry duplicates the first, after its INSERT has succeeded. duplicated := CreateInstanceModelInfo{ModelName: "gpt-duplicate", ModelTypes: []string{"chat"}, MaxTokens: 8192} code, err := NewModelProviderService().CreateProviderInstance( serviceTestContext(t, "user-1", "user-1"), @@ -606,6 +611,24 @@ func TestCreateProviderInstanceRollsBackInstanceWhenModelCreationFails(t *testin if modelCount != 0 { t.Fatalf("models = %d, want 0 after rollback", modelCount) } + code, err = NewModelProviderService().CreateProviderInstance( + serviceTestContext(t, "user-1", "user-1"), "OpenAI", "broken-instance", "sk-rollback", "", "default", "user-1", + []CreateInstanceModelInfo{duplicated, {ModelName: "gpt-second", ModelTypes: []string{"chat"}}}, + ) + if code != common.CodeSuccess || err != nil { + t.Fatalf("retry CreateProviderInstance() = (%v, %v)", code, err) + } + var instance entity.TenantModelInstance + if err := db.Where("instance_name = ?", "broken-instance").Take(&instance).Error; err != nil { + t.Fatal(err) + } + if err := db.Model(&entity.TenantModel{}).Where("instance_id = ?", instance.ID).Count(&modelCount).Error; err != nil { + t.Fatal(err) + } + if modelCount != 2 { + t.Fatalf("retry models = %d, want 2", modelCount) + } + } func TestModelProviderServiceAlterModelStatusByID(t *testing.T) { @@ -855,6 +878,11 @@ func TestResolveSavedProviderModelConfigUsePermissionCannotOverrideSavedEndpoint }).Error; err != nil { t.Fatalf("seed tenant member: %v", err) } + if err := db.Create(&entity.UserTenant{ + ID: "user-tenant-owner", UserID: "user-1", TenantID: "tenant-1", Role: string(permission.RoleOwner), Status: &activeStatus, + }).Error; err != nil { + t.Fatalf("seed tenant owner: %v", err) + } service := NewModelProviderService() ctx := serviceTestContext(t, "member-1", "tenant-1") config, err := service.ResolveSavedProviderModelConfig(ctx, "OpenAI", "default", "member-1", SavedProviderModelOptions{}) @@ -1017,7 +1045,7 @@ func TestModelFactoryResolveModelConfigByID(t *testing.T) { seedModelProviderServiceScope(t, db) ctx := serviceTestContext(t, "user-1", "user-1") - target, err := NewModelFactory().resolveConfig(ctx, ModelAccess{TenantID: "user-1", UserID: "user-1"}, entity.ModelTypeChat, "model-1") + target, err := NewModelFactory().resolveConfig(ctx, Identity{TenantID: "user-1", UserID: "user-1"}, entity.ModelTypeChat, "model-1") if err != nil { t.Fatalf("resolveConfig() error = %v", err) } @@ -1050,7 +1078,7 @@ func TestModelProviderServiceMissingProviderAndInstancePreserveLookupError(t *te } { t.Run(modelRef, func(t *testing.T) { _, err := NewModelFactory().resolveConfig( - serviceTestContext(t, "user-1", "tenant-1"), ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeEmbedding, modelRef, + serviceTestContext(t, "user-1", "tenant-1"), Identity{TenantID: "tenant-1"}, entity.ModelTypeEmbedding, modelRef, ) if !errors.Is(err, errModelConfigUnavailable) || !errors.Is(err, gorm.ErrRecordNotFound) { t.Fatalf("resolveConfig() error = %v, want unavailable record-not-found error", err) @@ -1077,13 +1105,13 @@ func TestModelProviderServiceDecodesCompatibleInstanceExtra(t *testing.T) { { name: "model ID", resolve: func() (*modelTarget, error) { - return solver.resolveConfig(serviceTestContext(t, "user-1", "user-1"), ModelAccess{TenantID: "user-1", UserID: "user-1"}, entity.ModelTypeChat, "model-1") + return solver.resolveConfig(serviceTestContext(t, "user-1", "user-1"), Identity{TenantID: "user-1", UserID: "user-1"}, entity.ModelTypeChat, "model-1") }, }, { name: "provider instance", resolve: func() (*modelTarget, error) { - return solver.resolveConfig(serviceTestContext(t, "user-1", "user-1"), ModelAccess{TenantID: "user-1"}, entity.ModelTypeChat, "gpt-test@default@OpenAI") + return solver.resolveConfig(serviceTestContext(t, "user-1", "user-1"), Identity{TenantID: "user-1"}, entity.ModelTypeChat, "gpt-test@default@OpenAI") }, }, } @@ -1358,7 +1386,7 @@ func TestModelFactoryResolveModelConfigReportsToolSupport(t *testing.T) { ctx := serviceTestContext(t, "user-1", "tenant-1") // A reference that passes validation still reads the persisted flag. - target, err := solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-active") + target, err := solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-active") if err != nil || target == nil || !target.SupportsTools { t.Fatalf("active chat model = (%v, %v), want tool support", target, err) } @@ -1371,7 +1399,7 @@ func TestModelFactoryResolveModelConfigReportsToolSupport(t *testing.T) { {"not enrolled as chat", "model-embedding"}, {"provider owned by another tenant", "model-foreign"}, } { - target, err := solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, tc.modelID) + target, err := solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, tc.modelID) if err == nil { t.Errorf("%s: err = nil, want a validation error", tc.name) } @@ -1432,7 +1460,7 @@ func TestModelFactoryResolveModelConfigPrefersTenantToolFlag(t *testing.T) { solver := NewModelFactory() ctx := serviceTestContext(t, "user-1", "tenant-1") - target, err := solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, toolsOn+"@default@OpenAI") + target, err := solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, toolsOn+"@default@OpenAI") if err != nil { t.Fatalf("tool-capable model: err = %v", err) } @@ -1440,7 +1468,7 @@ func TestModelFactoryResolveModelConfigPrefersTenantToolFlag(t *testing.T) { t.Errorf("tool support = true, want false: the tenant_model flag must win over the catalog") } - target, err = solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, toolsOff+"@default@OpenAI") + target, err = solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, toolsOff+"@default@OpenAI") if err != nil { t.Fatalf("tool-incapable model: err = %v", err) } @@ -1448,7 +1476,7 @@ func TestModelFactoryResolveModelConfigPrefersTenantToolFlag(t *testing.T) { t.Errorf("tool support = false, want true: the tenant_model flag must win over the catalog") } - target, err = solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, toolsOn+"-disabled@default@OpenAI") + target, err = solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, toolsOn+"-disabled@default@OpenAI") if err == nil || (target != nil && target.SupportsTools) { t.Errorf("disabled model = (%v, %v), want (false, error)", target, err) } @@ -1474,7 +1502,7 @@ func TestModelFactoryResolveModelConfigPropagatesLookupFailure(t *testing.T) { solver := NewModelFactory() ctx := serviceTestContext(t, "user-1", "tenant-1") - if target, err := solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-1"); err != nil || target == nil || !target.SupportsTools { + if target, err := solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-1"); err != nil || target == nil || !target.SupportsTools { t.Fatalf("baseline = (%v, %v), want tool support", target, err) } @@ -1485,11 +1513,11 @@ func TestModelFactoryResolveModelConfigPropagatesLookupFailure(t *testing.T) { t.Fatalf("failed to drop provider table: %v", err) } // Composite refs resolve the provider row directly. - if target, err := solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, "gpt-test@default@OpenAI"); err == nil { + if target, err := solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, "gpt-test@default@OpenAI"); err == nil { t.Errorf("composite provider lookup failure = (%v, nil), want a propagated error", target) } // The UUID path too. - if target, err := solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-1"); err == nil { + if target, err := solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-1"); err == nil { t.Errorf("uuid provider lookup failure = (%v, nil), want a propagated error", target) } @@ -1497,7 +1525,7 @@ func TestModelFactoryResolveModelConfigPropagatesLookupFailure(t *testing.T) { if err := db.Migrator().DropTable(&entity.TenantModel{}); err != nil { t.Fatalf("failed to drop model table: %v", err) } - if target, err := solver.resolveConfig(ctx, ModelAccess{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-1"); err == nil { + if target, err := solver.resolveConfig(ctx, Identity{TenantID: "tenant-1"}, entity.ModelTypeChat, "model-1"); err == nil { t.Errorf("model lookup failure = (%v, nil), want a propagated error", target) } } diff --git a/internal/service/model_service_tts_test.go b/internal/service/model_service_tts_test.go index aaf77019b7..98d4d4e069 100644 --- a/internal/service/model_service_tts_test.go +++ b/internal/service/model_service_tts_test.go @@ -58,7 +58,8 @@ func TestAudioSpeech_NoSelectionNoDefaultTTS(t *testing.T) { svc := NewModelCallService() text := "hello" - resp, code, err := svc.AudioSpeech(t.Context(), "", "tenant-1", &text, nil) + ctx := serviceTestContext(t, "tenant-1", "tenant-1") + resp, code, err := svc.AudioSpeech(ctx, "", &text, nil) if err == nil { t.Fatalf("expected error for missing default TTS model, got resp=%v", resp) } @@ -81,7 +82,7 @@ func TestResolveDefaultModelConfig_NilTTSPointer(t *testing.T) { pushServiceDB(t, testDB) insertTTSTestTenant(t, "tenant-2", nil) - _, err := NewModelFactory().resolveDefaultConfig(t.Context(), ModelAccess{TenantID: "tenant-2"}, entity.ModelTypeTTS) + _, err := NewModelFactory().resolveDefaultConfig(t.Context(), Identity{TenantID: "tenant-2"}, entity.ModelTypeTTS) if err == nil { t.Fatal("expected error for nil tts_id, got nil") } diff --git a/internal/service/model_stored_provider_name_test.go b/internal/service/model_stored_provider_name_test.go index a5e5f7933c..9f5c15a811 100644 --- a/internal/service/model_stored_provider_name_test.go +++ b/internal/service/model_stored_provider_name_test.go @@ -58,7 +58,7 @@ func TestModelFactoryResolvesStoredMistralOCRProvider(t *testing.T) { useModelProviderServiceTestDB(t, db) seedStoredMistralOCRProvider(t, db) - access := ModelAccess{TenantID: "user-1", UserID: "user-1"} + access := Identity{TenantID: "user-1", UserID: "user-1"} for _, tc := range []struct { name string modelRef string @@ -96,7 +96,7 @@ func TestModelFactoryResolveModelTypeForStoredMistralOCRProvider(t *testing.T) { useModelProviderServiceTestDB(t, db) seedStoredMistralOCRProvider(t, db) - types, err := NewModelFactory().ResolveModelType(serviceTestContext(t, "user-1", "user-1"), ModelAccess{TenantID: "user-1", UserID: "user-1"}, "mistral-ocr-2512@default@Mistral OCR") + types, err := NewModelFactory().ResolveModelType(serviceTestContext(t, "user-1", "user-1"), Identity{TenantID: "user-1", UserID: "user-1"}, "mistral-ocr-2512@default@Mistral OCR") if err != nil { t.Fatalf("ResolveModelType: %v", err) } @@ -115,28 +115,27 @@ func TestCatalogToolSupportForStoredMistralOCRProvider(t *testing.T) { } } -// The model call service resolves a "Mistral OCR" model by tenant model id and -// by composite reference. -func TestModelCallServiceResolvesStoredMistralOCRProvider(t *testing.T) { +// ModelFactory preserves the stored OCR capability for the Mistral OCR provider. +func TestModelFactoryResolvesStoredMistralOCRCapability(t *testing.T) { db := setupModelProviderServiceTestDB(t) useModelProviderServiceTestDB(t, db) seedStoredMistralOCRProvider(t, db) - svc := NewModelCallService() + factory := NewModelFactory() for _, modelRef := range []string{"model-mistral-ocr", "mistral-ocr-latest@default@Mistral OCR"} { t.Run(modelRef, func(t *testing.T) { - info, err := svc.resolveModelInfo(serviceTestContext(t, "user-1", "user-1"), modelRef, "user-1") + target, err := factory.resolveConfig(serviceTestContext(t, "user-1", "user-1"), Identity{TenantID: "user-1", UserID: "user-1"}, 0, modelRef) if err != nil { - t.Fatalf("resolveModelInfo(%q): %v", modelRef, err) + t.Fatalf("resolve model %q: %v", modelRef, err) } - if info.ProviderInfo == nil || info.ProviderInfo.Name != "Mistral" { - t.Fatalf("provider info = %+v, want the Mistral catalog entry", info.ProviderInfo) + if target.ProviderName != "Mistral OCR" { + t.Fatalf("provider name = %q, want the stored provider row", target.ProviderName) } - if info.ModelInfo == nil || info.ModelInfo.Name != "mistral-ocr-latest" { - t.Fatalf("model info = %+v, want mistral-ocr-latest", info.ModelInfo) + if target.ModelInfo == nil || target.ModelInfo.Name != "mistral-ocr-latest" { + t.Fatalf("model info = %+v, want mistral-ocr-latest", target.ModelInfo) } - if info.ProviderEntity == nil || info.ProviderEntity.ProviderName != "Mistral OCR" { - t.Fatalf("provider entity = %+v, want the stored row", info.ProviderEntity) + if !target.ModelInfo.ModelTypeMap["ocr"] || target.ModelInfo.ModelTypeMap["doc_parse"] { + t.Fatalf("model info = %+v, want OCR capability without unsupported document-parse capability", target.ModelInfo) } }) } diff --git a/internal/service/nav_embedder.go b/internal/service/nav_embedder.go index 80403e1be5..0562794a43 100644 --- a/internal/service/nav_embedder.go +++ b/internal/service/nav_embedder.go @@ -59,7 +59,7 @@ func (e *NavEmbedder) encode(ctx context.Context, tenantID string, texts []strin } name := e.embdModelName var model *modelModule.EmbeddingModel - access := ModelAccess{TenantID: tenantID} + access := Identity{TenantID: tenantID} var err error if name == "" { model, err = e.modelFactory.NewDefaultEmbeddingModel(ctx, access) diff --git a/internal/service/nav_merge_llm.go b/internal/service/nav_merge_llm.go index 0594fb12a0..f490de365e 100644 --- a/internal/service/nav_merge_llm.go +++ b/internal/service/nav_merge_llm.go @@ -225,7 +225,7 @@ func navChatText(ctx context.Context, modelFactory *ModelFactory, tenantID, mode if modelFactory == nil { return "", fmt.Errorf("datasetnav: model factory not initialized") } - access := ModelAccess{TenantID: tenantID} + access := Identity{TenantID: tenantID} var ( chatModel *modelModule.ChatModel err error diff --git a/internal/service/openai_chat.go b/internal/service/openai_chat.go index 3ef0e733f7..50f203d0b6 100644 --- a/internal/service/openai_chat.go +++ b/internal/service/openai_chat.go @@ -110,16 +110,14 @@ type OpenAIStreamEvent struct { // keeping handler-level concerns (message filtering, generation config merge, // reference metadata enrichment) on the service itself. type OpenAIChatService struct { - chatSvc *ChatService - tenantLLMSvc *TenantLLMService - pipeline *ChatPipelineService + chatSvc *ChatService + pipeline *ChatPipelineService } func NewOpenAIChatService() *OpenAIChatService { return &OpenAIChatService{ - chatSvc: NewChatService(), - tenantLLMSvc: NewTenantLLMService(), - pipeline: NewChatPipelineService(), + chatSvc: NewChatService(), + pipeline: NewChatPipelineService(), } } @@ -263,15 +261,10 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID } } if req.Model != "model" { - if _, mErr := s.pipeline.ModelFactory.ResolveInfo(ctx, ModelAccess{UserID: userID, TenantID: dialog.TenantID}, entity.ModelTypeChat, resolvedModel); mErr != nil { + if _, mErr := s.pipeline.ModelFactory.ResolveInfo(ctx, Identity{UserID: userID, TenantID: dialog.TenantID}, entity.ModelTypeChat, resolvedModel); mErr != nil { s.writeArgError(c, fmt.Sprintf("`llm_id` %s doesn't exist", req.Model)) return } - apiKey, apiErr := s.tenantLLMSvc.GetAPIKeyFromInstance(ctx, dialog.TenantID, req.Model) - if apiErr != nil || apiKey == "" { - s.writeDataError(c, fmt.Sprintf("Cannot use specified model %s.", req.Model)) - return - } dialog.LLMID = resolvedModel } diff --git a/internal/service/related_question.go b/internal/service/related_question.go index 973f7f6eb0..0f55c82c6a 100644 --- a/internal/service/related_question.go +++ b/internal/service/related_question.go @@ -19,6 +19,7 @@ package service import ( "context" "encoding/json" + "errors" "fmt" "regexp" "strings" @@ -29,18 +30,19 @@ import ( ) // GenerateRelatedQuestions generates related search questions for chat/searchbot endpoints. -func GenerateRelatedQuestions(ctx context.Context, userID, question, searchID string, searchSvc *SearchService, tenantSvc *TenantService, modelFactory *ModelFactory) ([]string, error) { +func GenerateRelatedQuestions(ctx context.Context, question, searchID string, searchSvc *SearchService, modelFactory *ModelFactory) ([]string, error) { if modelFactory == nil { return nil, fmt.Errorf("model factory not configured") } - searchConfig, err := relatedQuestionsSearchConfig(ctx, userID, searchID, searchSvc) + searchConfig, err := relatedQuestionsSearchConfig(ctx, searchID, searchSvc) if err != nil { return nil, err } - modelID := relatedQuestionsModelID(ctx, userID, searchConfig, tenantSvc) - if modelID == "" { - return nil, fmt.Errorf("no chat model configured") + access, err := ResolveIdentity(ctx) + if err != nil { + return nil, err } + modelID, _ := searchConfig["chat_id"].(string) prompt, err := (prompts.EmbeddedPromptLoader{}).Load("related_question") if err != nil { return nil, err @@ -49,7 +51,15 @@ func GenerateRelatedQuestions(ctx context.Context, userID, question, searchID st {Role: "system", Content: prompt}, {Role: "user", Content: "\nKeywords: " + question + "\nRelated search terms:\n "}, } - chatModel, err := modelFactory.NewChatModel(ctx, ModelAccess{TenantID: userID}, modelID) + var chatModel *modelModule.ChatModel + if modelID == "" { + chatModel, err = modelFactory.NewDefaultChatModel(ctx, access) + if errors.Is(err, ErrNoDefaultModel) { + return nil, fmt.Errorf("no chat model configured") + } + } else { + chatModel, err = modelFactory.NewChatModel(ctx, access, modelID) + } if err != nil { return nil, err } @@ -63,14 +73,14 @@ func GenerateRelatedQuestions(ctx context.Context, userID, question, searchID st return []string{}, nil } -func relatedQuestionsSearchConfig(ctx context.Context, userID, searchID string, searchSvc *SearchService) (map[string]interface{}, error) { +func relatedQuestionsSearchConfig(ctx context.Context, searchID string, searchSvc *SearchService) (map[string]interface{}, error) { if searchID == "" { return map[string]interface{}{}, nil } if searchSvc == nil { return nil, fmt.Errorf("search service not configured") } - detail, err := searchSvc.GetDetail(ctx, userID, searchID) + detail, err := searchSvc.GetDetail(ctx, searchID) if err != nil { return nil, err } @@ -87,18 +97,6 @@ func relatedQuestionsSearchConfigFromDetail(detail map[string]interface{}) map[s return map[string]interface{}{} } -func relatedQuestionsModelID(ctx context.Context, tenantID string, searchConfig map[string]interface{}, tenantSvc *TenantService) string { - modelID, _ := searchConfig["chat_id"].(string) - if modelID != "" || tenantSvc == nil { - return modelID - } - defaultModel, err := tenantSvc.GetDefaultModelNameForTenant(ctx, tenantID, entity.ModelTypeChat) - if err == nil { - modelID = defaultModel - } - return modelID -} - func relatedQuestionsConfig(searchConfig map[string]interface{}) *modelModule.ChatConfig { var genConf map[string]interface{} switch v := searchConfig["llm_setting"].(type) { diff --git a/internal/service/search.go b/internal/service/search.go index 9385aef75e..e8bc2ed2da 100644 --- a/internal/service/search.go +++ b/internal/service/search.go @@ -31,23 +31,15 @@ import ( // SearchService search service type SearchService struct { - searchDAO *dao.SearchDAO - datasetDAO *dao.KnowledgebaseDAO - tenantService *TenantService + searchDAO *dao.SearchDAO + datasetDAO *dao.KnowledgebaseDAO } // NewSearchService create search service func NewSearchService() *SearchService { return &SearchService{ - searchDAO: dao.NewSearchDAO(), - datasetDAO: dao.NewKnowledgebaseDAO(), - tenantService: NewTenantService(), - } -} - -func (s *SearchService) SetTenantService(tenantService *TenantService) { - if tenantService != nil { - s.tenantService = tenantService + searchDAO: dao.NewSearchDAO(), + datasetDAO: dao.NewKnowledgebaseDAO(), } } @@ -81,11 +73,18 @@ type SearchShareDetail struct { } // ListSearches list search apps with advanced filtering (equivalent to list_search_app) -func (s *SearchService) ListSearches(ctx context.Context, userID string, keywords string, page, pageSize int, terms []dao.OrderTerm, ownerIDs []string) (*ListSearchAppsResponse, error) { - resourceIDs, err := accessibleSearchAppIDs(ctx, userID, permission.OperationRead) +func (s *SearchService) ListSearches(ctx context.Context, keywords string, page, pageSize int, terms []dao.OrderTerm, ownerIDs []string) (*ListSearchAppsResponse, error) { + scope, err := permission.NewDatabaseChecker(dao.DB).Scope(ctx, permission.Subject{}, permission.ScopeQuery{ + Kind: permission.ResourceKindSearchApp, + Operation: permission.OperationRead, + }) if err != nil { return nil, err } + var resourceIDs []string + if scope.Mode == permission.ScopeIDs { + resourceIDs = scope.ResourceIDs + } searches, total, err := s.searchDAO.ListByResourceIDs(ctx, dao.DB, resourceIDs, ownerIDs, page, pageSize, terms, keywords) if err != nil { return nil, err @@ -95,7 +94,7 @@ func (s *SearchService) ListSearches(ctx context.Context, userID string, keyword searchApps := make([]map[string]interface{}, len(searches)) checker := permission.NewDatabaseChecker(dao.DB) for i, search := range searches { - access, err := checker.ResolveAccess(ctx, permission.Subject{UserID: userID}, permission.ResourceRef{ + access, err := checker.ResolveAccess(ctx, permission.Subject{}, permission.ResourceRef{ Kind: permission.ResourceKindSearchApp, ID: search.ID, }) @@ -153,12 +152,12 @@ type CreateSearchResponse struct { // 5. Set fields: id, name, description, tenant_id, created_by // 6. Save to database within DB.atomic() transaction // 7. Return {search_id: id} on success -func (s *SearchService) CreateSearch(ctx context.Context, userID string, name string, description *string) (*CreateSearchResponse, error) { - tenantID, err := requireCurrentTenantID(ctx, userID) +func (s *SearchService) CreateSearch(ctx context.Context, name string, description *string) (*CreateSearchResponse, error) { + identity, err := ResolveIdentity(ctx) if err != nil { return nil, err } - if err := checkFeatureAction(ctx, userID, "search", permission.ActionWrite); err != nil { + if err := permission.NewDatabaseChecker(dao.DB).CheckFeature(ctx, permission.Subject{UserID: identity.UserID}, "search", permission.ActionWrite); err != nil { return nil, err } if err := common.ValidateName(name); err != nil { @@ -171,7 +170,7 @@ func (s *SearchService) CreateSearch(ctx context.Context, userID string, name st // Generate unique name (same as Python duplicate_name) // search.name is a 128-byte column; keep generated names within it. uniqueName, err := common.UniqueName(name, 128, func(candidate string) (bool, error) { - existing, err := s.searchDAO.GetByNameAndTenant(ctx, dao.DB, candidate, tenantID) + existing, err := s.searchDAO.GetByNameAndTenant(ctx, dao.DB, candidate, identity.TenantID) if err != nil { return false, err } @@ -185,9 +184,9 @@ func (s *SearchService) CreateSearch(ctx context.Context, userID string, name st // Create search entity search := &entity.Search{ ID: searchID, - TenantID: tenantID, + TenantID: identity.TenantID, Name: uniqueName, - CreatedBy: userID, + CreatedBy: identity.UserID, SearchConfig: make(entity.JSONMap), } @@ -204,7 +203,7 @@ func (s *SearchService) CreateSearch(ctx context.Context, userID string, name st if err := s.searchDAO.Create(ctx, tx, search); err != nil { return err } - return grantCreatedResourceOwner(ctx, tx, userID, tenantID, permission.ResourceKindSearchApp, search.ID) + return grantCreatedResourceOwner(ctx, tx, identity.UserID, identity.TenantID, permission.ResourceKindSearchApp, search.ID) }); err != nil { return nil, fmt.Errorf("failed to create search: %w", err) } @@ -214,8 +213,11 @@ func (s *SearchService) CreateSearch(ctx context.Context, userID string, name st }, nil } -func (s *SearchService) GetSearchDetail(ctx context.Context, userID string, searchID string) (*entity.Search, error) { - if err := checkSearchAppPermission(ctx, userID, searchID, permission.OperationRead); err != nil { +func (s *SearchService) GetSearchDetail(ctx context.Context, searchID string) (*entity.Search, error) { + if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.Subject{}, permission.ResourceRef{ + Kind: permission.ResourceKindSearchApp, + ID: searchID, + }, permission.OperationRead); err != nil { return nil, err } search, err := s.searchDAO.GetByID(ctx, dao.DB, searchID) @@ -228,8 +230,8 @@ func (s *SearchService) GetSearchDetail(ctx context.Context, userID string, sear // GetSearchShareDetail returns the joined share-detail payload for public // searchbot pages after verifying the caller can access the search app. -func (s *SearchService) GetSearchShareDetail(ctx context.Context, userID, searchID string) (*SearchShareDetail, error) { - if _, err := s.GetSearchDetail(ctx, userID, searchID); err != nil { +func (s *SearchService) GetSearchShareDetail(ctx context.Context, searchID string) (*SearchShareDetail, error) { + if _, err := s.GetSearchDetail(ctx, searchID); err != nil { return nil, err } @@ -254,8 +256,11 @@ func (s *SearchService) GetSearchShareDetail(ctx context.Context, userID, search } // DeleteSearch deletes a search app by ID -func (s *SearchService) DeleteSearch(ctx context.Context, userID string, searchID string) error { - if err := checkSearchAppPermission(ctx, userID, searchID, permission.OperationDelete); err != nil { +func (s *SearchService) DeleteSearch(ctx context.Context, searchID string) error { + if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.Subject{}, permission.ResourceRef{ + Kind: permission.ResourceKindSearchApp, + ID: searchID, + }, permission.OperationDelete); err != nil { return err } if err := s.searchDAO.DeleteByID(ctx, dao.DB, searchID); err != nil { @@ -266,7 +271,7 @@ func (s *SearchService) DeleteSearch(ctx context.Context, userID string, searchI } type SearchCompletionPlan struct { - UserID string + Identity Identity SearchID string Question string DatasetIDs []string @@ -274,11 +279,7 @@ type SearchCompletionPlan struct { Options AskStreamOptions } -func (s *SearchService) PrepareCompletion(ctx context.Context, userID, searchID string, req *SearchCompletionsRequest) (*SearchCompletionPlan, common.ErrorCode, error) { - userID = strings.TrimSpace(userID) - if userID == "" { - return nil, common.CodeBadRequest, fmt.Errorf("user id is required") - } +func (s *SearchService) PrepareCompletion(ctx context.Context, searchID string, req *SearchCompletionsRequest) (*SearchCompletionPlan, common.ErrorCode, error) { searchID = strings.TrimSpace(searchID) if searchID == "" { return nil, common.CodeBadRequest, fmt.Errorf("search_id is required") @@ -295,7 +296,7 @@ func (s *SearchService) PrepareCompletion(ctx context.Context, userID, searchID return nil, common.CodeArgumentError, fmt.Errorf("question is required") } - searchDetail, err := s.GetDetail(ctx, userID, searchID) + searchDetail, err := s.GetDetail(ctx, searchID) if err != nil || searchDetail == nil { if permissionresponse.IsPermissionError(err) { code, permissionErr := permissionresponse.NormalizeHidden(err) @@ -314,27 +315,29 @@ func (s *SearchService) PrepareCompletion(ctx context.Context, userID, searchID } for _, datasetID := range datasetIDs { - if err := CheckDatasetAccess(ctx, permission.Subject{UserID: userID}, datasetID, permission.OperationUse); err != nil { + if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.Subject{}, permission.ResourceRef{ + Kind: permission.ResourceKindDataset, + ID: datasetID, + }, permission.OperationUse); err != nil { code, permissionErr := permissionresponse.Normalize(err) return nil, code, permissionErr } } + identity, err := ResolveIdentity(ctx) + if err != nil { + code, normalizedErr := permissionresponse.Normalize(err) + return nil, code, normalizedErr + } modelID, _ := stringFromSearchConfig(searchConfig["chat_id"]) if modelID == "" { - tenantSvc := s.tenantService - if tenantSvc == nil { - tenantSvc = NewTenantService() - } - var defaultModelName string - defaultModelName, err = tenantSvc.GetDefaultModelNameForUser(ctx, userID, entity.ModelTypeChat) - if err == nil { - modelID = strings.TrimSpace(defaultModelName) + if info, infoErr := NewModelFactory().ResolveDefaultInfo(ctx, identity, entity.ModelTypeChat); infoErr == nil { + modelID = ModelReference(info) } } return &SearchCompletionPlan{ - UserID: userID, + Identity: identity, SearchID: searchID, Question: question, DatasetIDs: datasetIDs, @@ -518,8 +521,11 @@ type UpdateSearchRequest struct { Avatar *string `json:"avatar,omitempty"` } -func (s *SearchService) UpdateSearch(ctx context.Context, userID string, searchID string, req *UpdateSearchRequest) (*entity.Search, error) { - if err := checkSearchAppPermission(ctx, userID, searchID, permission.OperationUpdate); err != nil { +func (s *SearchService) UpdateSearch(ctx context.Context, searchID string, req *UpdateSearchRequest) (*entity.Search, error) { + if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.Subject{}, permission.ResourceRef{ + Kind: permission.ResourceKindSearchApp, + ID: searchID, + }, permission.OperationUpdate); err != nil { return nil, err } @@ -589,8 +595,11 @@ func (s *SearchService) UpdateSearch(ctx context.Context, userID string, searchI } // GetDetail returns Search App configuration to a caller that is executing it. -func (s *SearchService) GetDetail(ctx context.Context, userID, searchID string) (map[string]interface{}, error) { - if err := checkSearchAppPermission(ctx, userID, searchID, permission.OperationRun); err != nil { +func (s *SearchService) GetDetail(ctx context.Context, searchID string) (map[string]interface{}, error) { + if err := permission.NewDatabaseChecker(dao.DB).CheckResource(ctx, permission.Subject{}, permission.ResourceRef{ + Kind: permission.ResourceKindSearchApp, + ID: searchID, + }, permission.OperationRun); err != nil { return nil, err } search, err := s.searchDAO.GetByID(ctx, dao.DB, searchID) @@ -609,7 +618,7 @@ func (s *SearchService) GetDetail(ctx context.Context, userID, searchID string) "update_time": search.UpdateTime, "search_config": map[string]interface{}(search.SearchConfig), } - access, err := permission.NewDatabaseChecker(dao.DB).ResolveAccess(ctx, permission.Subject{UserID: userID}, permission.ResourceRef{ + access, err := permission.NewDatabaseChecker(dao.DB).ResolveAccess(ctx, permission.Subject{}, permission.ResourceRef{ Kind: permission.ResourceKindSearchApp, ID: search.ID, }) diff --git a/internal/service/search_permission.go b/internal/service/search_permission.go deleted file mode 100644 index b468a02c9a..0000000000 --- a/internal/service/search_permission.go +++ /dev/null @@ -1,47 +0,0 @@ -// -// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -package service - -import ( - "context" - - "ragflow/internal/dao" - "ragflow/internal/permission" -) - -func checkSearchAppPermission(ctx context.Context, userID, searchID string, operation permission.Operation) error { - return permission.NewDatabaseChecker(dao.DB).CheckResource( - ctx, - permission.Subject{UserID: userID}, - permission.ResourceRef{Kind: permission.ResourceKindSearchApp, ID: searchID}, - operation, - ) -} - -func accessibleSearchAppIDs(ctx context.Context, userID string, operation permission.Operation) ([]string, error) { - scope, err := permission.NewDatabaseChecker(dao.DB).Scope(ctx, permission.Subject{UserID: userID}, permission.ScopeQuery{ - Kind: permission.ResourceKindSearchApp, - Operation: operation, - }) - if err != nil { - return nil, err - } - if scope.Mode != permission.ScopeIDs { - return []string{}, nil - } - return scope.ResourceIDs, nil -} diff --git a/internal/service/search_share_detail_test.go b/internal/service/search_share_detail_test.go index b1f899cac7..c94a5cfec5 100644 --- a/internal/service/search_share_detail_test.go +++ b/internal/service/search_share_detail_test.go @@ -79,7 +79,7 @@ func TestSearchServiceGetSearchShareDetail(t *testing.T) { } ctx := serviceTestContext(t, "tenant-1", "tenant-1") - detail, err := NewSearchService().GetSearchShareDetail(ctx, "tenant-1", "search-1") + detail, err := NewSearchService().GetSearchShareDetail(ctx, "search-1") if err != nil { t.Fatalf("GetSearchShareDetail failed: %v", err) } @@ -147,7 +147,7 @@ func TestSearchServiceGetSearchShareDetailRejectsUnauthorizedUser(t *testing.T) } ctx := serviceTestContext(t, "user-2", "tenant-1", "tenant-1") - _, err := NewSearchService().GetSearchShareDetail(ctx, "user-2", "search-1") + _, err := NewSearchService().GetSearchShareDetail(ctx, "search-1") if err == nil { t.Fatal("expected permission error") } diff --git a/internal/service/search_test.go b/internal/service/search_test.go index fbab4745b2..5c1d2d0c18 100644 --- a/internal/service/search_test.go +++ b/internal/service/search_test.go @@ -23,6 +23,7 @@ import ( "strings" "testing" + "ragflow/internal/common" "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/permission" @@ -57,7 +58,7 @@ func TestSearchServiceCreateRejectsEmptyName(t *testing.T) { setupSearchServiceTestDB(t) ctx := serviceTestContext(t, "tenant-1", "tenant-1") - _, err := NewSearchService().CreateSearch(ctx, "tenant-1", " ", nil) + _, err := NewSearchService().CreateSearch(ctx, " ", nil) if err == nil { t.Fatal("expected empty name validation error") } @@ -74,7 +75,7 @@ func TestSearchServiceUpdateRejectsUnauthorizedSearchID(t *testing.T) { Name: "New Name", SearchConfig: map[string]interface{}{}, } - _, err := NewSearchService().UpdateSearch(ctx, "user-2", "invalid_search_id", req) + _, err := NewSearchService().UpdateSearch(ctx, "invalid_search_id", req) if err == nil { t.Fatal("expected authorization error") } @@ -86,7 +87,7 @@ func TestSearchServiceUpdateRejectsUnauthorizedSearchID(t *testing.T) { func TestSearchServiceCreateAndUpdateRoundTrip(t *testing.T) { setupSearchServiceTestDB(t) ctx := serviceTestContext(t, "tenant-1", "tenant-1", "user-2") - created, err := NewSearchService().CreateSearch(ctx, "tenant-1", "My Search", nil) + created, err := NewSearchService().CreateSearch(ctx, "My Search", nil) if err != nil { t.Fatalf("CreateSearch failed: %v", err) } @@ -99,7 +100,7 @@ func TestSearchServiceCreateAndUpdateRoundTrip(t *testing.T) { Name: "Hijacked Name", SearchConfig: map[string]interface{}{}, } - _, err = NewSearchService().UpdateSearch(serviceTestContext(t, "user-2", "tenant-1", "tenant-1"), "user-2", created.SearchID, req) + _, err = NewSearchService().UpdateSearch(serviceTestContext(t, "user-2", "tenant-1", "tenant-1"), created.SearchID, req) if !errors.Is(err, permission.ErrPermissionDenied) { t.Fatalf("expected permission denied, got %v", err) } @@ -109,7 +110,7 @@ func TestSearchServiceCreateAndUpdateRoundTrip(t *testing.T) { Name: "Updated Name", SearchConfig: map[string]interface{}{"summary": true, "keywords_similarity_weight": 0.7}, } - updated, err := NewSearchService().UpdateSearch(ctx, "tenant-1", created.SearchID, req) + updated, err := NewSearchService().UpdateSearch(ctx, created.SearchID, req) if err != nil { t.Fatalf("owner UpdateSearch failed: %v", err) } @@ -134,6 +135,39 @@ func TestSearchServiceCreateAndUpdateRoundTrip(t *testing.T) { } } +func TestPrepareCompletionKeepsActiveTenantIdentity(t *testing.T) { + setupSearchServiceTestDB(t) + ctx := serviceTestContext(t, "user-1", "team-1") + status := string(entity.StatusValid) + if err := dao.DB.Create(&entity.UserTenant{ + ID: "team-owner", UserID: "team-1", TenantID: "team-1", Role: string(permission.RoleOwner), Status: &status, + }).Error; err != nil { + t.Fatalf("create tenant owner membership: %v", err) + } + if err := dao.DB.Create(&entity.Knowledgebase{ + ID: "kb-1", TenantID: "team-1", CreatedBy: "user-1", Name: "dataset", Permission: "team", Status: &status, + }).Error; err != nil { + t.Fatalf("create dataset: %v", err) + } + if err := dao.DB.Create(&entity.Search{ + ID: "search-1", TenantID: "team-1", CreatedBy: "user-1", Name: "search", + SearchConfig: entity.JSONMap{"kb_ids": []interface{}{"kb-1"}}, Status: &status, + }).Error; err != nil { + t.Fatalf("create search: %v", err) + } + + plan, code, err := NewSearchService().PrepareCompletion(ctx, "search-1", &SearchCompletionsRequest{Question: "question"}) + if err != nil { + t.Fatalf("PrepareCompletion() error = %v", err) + } + if code != common.CodeSuccess { + t.Fatalf("PrepareCompletion() code = %v, want success", code) + } + if plan.Identity != (Identity{UserID: "user-1", TenantID: "team-1"}) { + t.Fatalf("Identity = %+v, want user-1 in team-1", plan.Identity) + } +} + func TestBuildSearchConfigResponseOmitsVectorSimilarityWeight(t *testing.T) { searchConfig := map[string]interface{}{ "similarity_threshold": 0.2, @@ -165,7 +199,7 @@ func TestSearchServiceListSearchesReturnsOwnerDisplayFields(t *testing.T) { createSearchServiceTestSearch(t, "search-1", "user-1", "Search One") ctx := serviceTestContext(t, "user-1", "user-1") - result, err := NewSearchService().ListSearches(ctx, "user-1", "", 0, 0, []dao.OrderTerm{{Column: "create_time", Desc: true}}, nil) + result, err := NewSearchService().ListSearches(ctx, "", 0, 0, []dao.OrderTerm{{Column: "create_time", Desc: true}}, nil) if err != nil { t.Fatalf("ListSearches failed: %v", err) } @@ -183,7 +217,7 @@ func TestSearchServiceListSearchesNicknameFallsBackToTenantID(t *testing.T) { createSearchServiceTestSearch(t, "search-1", "user-1", "Search One") ctx := serviceTestContext(t, "user-1", "user-1") - result, err := NewSearchService().ListSearches(ctx, "user-1", "", 0, 0, []dao.OrderTerm{{Column: "create_time", Desc: true}}, nil) + result, err := NewSearchService().ListSearches(ctx, "", 0, 0, []dao.OrderTerm{{Column: "create_time", Desc: true}}, nil) if err != nil { t.Fatalf("ListSearches failed: %v", err) } @@ -219,7 +253,7 @@ func TestSearchServiceListSearchesUsesPermissionScope(t *testing.T) { } ctx := serviceTestContext(t, "member-1", "team-tenant", "team-owner") - result, err := NewSearchService().ListSearches(ctx, "member-1", "", 0, 0, nil, nil) + result, err := NewSearchService().ListSearches(ctx, "", 0, 0, nil, nil) if err != nil { t.Fatalf("ListSearches failed: %v", err) } @@ -262,10 +296,10 @@ func TestSearchServiceSeparatesSearchReadFromRunPermission(t *testing.T) { service := NewSearchService() ctx := serviceTestContext(t, "member-1", "team-tenant", "team-owner") - if _, err := service.GetSearchDetail(ctx, "member-1", search.ID); !errors.Is(err, permission.ErrPermissionDenied) { + if _, err := service.GetSearchDetail(ctx, search.ID); !errors.Is(err, permission.ErrPermissionDenied) { t.Fatalf("tenant member should not be able to read another user's Search App metadata, got %v", err) } - if _, err := service.GetDetail(ctx, "member-1", search.ID); !errors.Is(err, permission.ErrPermissionDenied) { + if _, err := service.GetDetail(ctx, search.ID); !errors.Is(err, permission.ErrPermissionDenied) { t.Fatalf("tenant member should not be able to run another user's Search App, got %v", err) } } @@ -303,7 +337,7 @@ func TestSearchServiceListSearchesPaginatesOwnerFilter(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - result, err := NewSearchService().ListSearches(ctx, "tenant-1", "", tt.page, tt.pageSize, []dao.OrderTerm{{Column: "create_time", Desc: true}}, []string{"tenant-1"}) + result, err := NewSearchService().ListSearches(ctx, "", tt.page, tt.pageSize, []dao.OrderTerm{{Column: "create_time", Desc: true}}, []string{"tenant-1"}) if err != nil { t.Fatalf("ListSearches failed: %v", err) } diff --git a/internal/service/skill_indexer.go b/internal/service/skill_indexer.go index 3cb484f03f..a63b3b4f26 100644 --- a/internal/service/skill_indexer.go +++ b/internal/service/skill_indexer.go @@ -955,7 +955,7 @@ func (s *SkillIndexerService) generateEmbedding(ctx context.Context, text, embdI return nil, fmt.Errorf("embedding model ID not configured") } - embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, ModelAccess{TenantID: tenantID}, embdID) + embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, Identity{TenantID: tenantID}, embdID) if err != nil { return nil, fmt.Errorf("failed to get embedding model: %w", err) } @@ -986,7 +986,7 @@ func (s *SkillIndexerService) generateEmbeddings(ctx context.Context, texts []st } common.Info(fmt.Sprintf("Getting embedding model for %s", embdID)) - embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, ModelAccess{TenantID: tenantID}, embdID) + embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, Identity{TenantID: tenantID}, embdID) if err != nil { common.Error(fmt.Sprintf("Failed to get embedding model: %v", err), err) return nil, fmt.Errorf("failed to get embedding model: %w", err) @@ -1021,7 +1021,7 @@ func (s *SkillIndexerService) getEmbeddingDimension(ctx context.Context, tenantI return 0, fmt.Errorf("embedding model ID not configured") } - embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, ModelAccess{TenantID: tenantID}, embdID) + embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, Identity{TenantID: tenantID}, embdID) if err != nil { return 0, fmt.Errorf("failed to get embedding model: %w", err) } diff --git a/internal/service/skill_search.go b/internal/service/skill_search.go index d0ad890de7..bdfd1e285b 100644 --- a/internal/service/skill_search.go +++ b/internal/service/skill_search.go @@ -660,7 +660,7 @@ func (s *SkillSearchService) getEmbedding(ctx context.Context, text, embdID, ten return nil, fmt.Errorf("embedding model ID not configured") } - embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, ModelAccess{TenantID: tenantID}, embdID) + embeddingModel, err := s.modelFactory.NewEmbeddingModel(ctx, Identity{TenantID: tenantID}, embdID) if err != nil { return nil, fmt.Errorf("failed to get embedding model: %w", err) } diff --git a/internal/service/skill_space.go b/internal/service/skill_space.go index 2bcaa86a53..ac3610443e 100644 --- a/internal/service/skill_space.go +++ b/internal/service/skill_space.go @@ -165,11 +165,13 @@ func (s *SkillSpaceService) CreateSpace(ctx context.Context, userID string, req } defaultEmbdID := req.EmbdID if defaultEmbdID == "" { - defaultEmbdID, _, err = NewTenantService().GetDefaultModelRefsForUser(ctx, userID, tenantID, entity.ModelTypeEmbedding) - if err != nil { - common.Error("Failed to resolve default embedding model for skill space", err, + info, modelErr := NewModelFactory().ResolveDefaultInfo(ctx, Identity{UserID: userID, TenantID: tenantID}, entity.ModelTypeEmbedding) + if modelErr == nil { + defaultEmbdID = ModelReference(info) + } else if !errors.Is(modelErr, ErrNoDefaultModel) { + common.Error("Failed to resolve default embedding model for skill space", modelErr, zap.String("tenantID", tenantID), zap.String("userID", userID)) - return nil, common.CodeOperatingError, fmt.Errorf("failed to resolve default embedding model: %w", err) + return nil, common.CodeOperatingError, fmt.Errorf("failed to resolve default embedding model: %w", modelErr) } } diff --git a/internal/service/stub.go b/internal/service/stub.go index 60fae2831f..ade604467f 100644 --- a/internal/service/stub.go +++ b/internal/service/stub.go @@ -74,9 +74,9 @@ func resolveDefaultTenantIDEE(context.Context, string) (string, bool, error) { return "", false, nil } -// resolveUserAuthTenant keeps the token's tenant unchanged. -func resolveUserAuthTenant(_ context.Context, _ *entity.User, tenantID string) (string, error) { - return tenantID, nil +// resolveUserAuthTenant discards tenant context because community credentials identify only a user. +func resolveUserAuthTenant(_ context.Context, _ *entity.User, _ string) (string, error) { + return "", nil } // LoginChannel is one authentication channel advertised to the login page. diff --git a/internal/service/tenant.go b/internal/service/tenant.go index 3e3370f469..f5f3414928 100644 --- a/internal/service/tenant.go +++ b/internal/service/tenant.go @@ -35,14 +35,10 @@ import ( // TenantService tenant service type TenantService struct { - tenantDAO *dao.TenantDAO - userTenantDAO *dao.UserTenantDAO - userDAO *dao.UserDAO - modelProviderDAO *dao.TenantModelProviderDAO - modelInstanceDAO *dao.TenantModelInstanceDAO - modelDAO *dao.TenantModelDAO - kbDAO *dao.KnowledgebaseDAO - docEngine engine.DocEngine + userTenantDAO *dao.UserTenantDAO + userDAO *dao.UserDAO + kbDAO *dao.KnowledgebaseDAO + docEngine engine.DocEngine } type tenantInviteRole struct { @@ -108,55 +104,13 @@ func (s *TenantService) checkTenantInviteEnabled(ctx context.Context, tenantID s // NewTenantService create tenant service func NewTenantService() *TenantService { return &TenantService{ - tenantDAO: dao.NewTenantDAO(), - userTenantDAO: dao.NewUserTenantDAO(), - userDAO: dao.NewUserDAO(), - modelProviderDAO: dao.NewTenantModelProviderDAO(), - modelInstanceDAO: dao.NewTenantModelInstanceDAO(), - modelDAO: dao.NewTenantModelDAO(), - kbDAO: dao.NewKnowledgebaseDAO(), - docEngine: engine.Get(), + userTenantDAO: dao.NewUserTenantDAO(), + userDAO: dao.NewUserDAO(), + kbDAO: dao.NewKnowledgebaseDAO(), + docEngine: engine.Get(), } } -// TenantInfoResponse tenant information response -type TenantInfoResponse struct { - TenantID string `json:"tenant_id"` - Name *string `json:"name,omitempty"` - LLMID string `json:"llm_id"` - EmbDID string `json:"embd_id"` - RerankID string `json:"rerank_id"` - ASRID string `json:"asr_id"` - Img2TxtID string `json:"img2txt_id"` - TTSID string `json:"tts_id"` - ParserIDs string `json:"parser_ids"` - Role string `json:"role"` -} - -// GetTenantInfo returns the user's own tenant and its owner membership settings. -func (s *TenantService) GetTenantInfo(ctx context.Context, userID string) (*TenantInfoResponse, error) { - membership, tenantID, err := s.currentUserTenant(ctx, userID) - if err != nil { - return nil, err - } - tenant, err := s.tenantDAO.GetByID(ctx, dao.DB, tenantID) - if err != nil { - return nil, err - } - return &TenantInfoResponse{ - TenantID: tenant.ID, - Name: tenant.Name, - LLMID: membership.LLMID, - EmbDID: membership.EmbdID, - RerankID: membership.RerankID, - ASRID: membership.ASRID, - Img2TxtID: membership.Img2TxtID, - TTSID: ptrStringValue(membership.TTSID), - ParserIDs: tenant.ParserIDs, - Role: membership.Role, - }, nil -} - // TenantListItem tenant list item response type TenantListItem struct { TenantID string `json:"tenant_id"` @@ -169,66 +123,6 @@ type TenantListItem struct { DeltaSeconds float64 `json:"delta_seconds"` } -// TenantLLMService tenant LLM service -// This service handles operations related to tenant-specific LLM configurations -type TenantLLMService struct { - modelProviderDAO *dao.TenantModelProviderDAO - modelInstanceDAO *dao.TenantModelInstanceDAO - modelDAO *dao.TenantModelDAO -} - -// NewTenantLLMService creates a new TenantLLMService instance -func NewTenantLLMService() *TenantLLMService { - return &TenantLLMService{ - modelProviderDAO: dao.NewTenantModelProviderDAO(), - modelInstanceDAO: dao.NewTenantModelInstanceDAO(), - modelDAO: dao.NewTenantModelDAO(), - } -} - -// GetAPIKeyFromInstance returns the API key for the given composite model name -// by looking it up in the tenant_model_instance table. compositeModelName is in -// "model@instance@provider" or "model@provider" format. -func (s *TenantLLMService) GetAPIKeyFromInstance(ctx context.Context, tenantID, compositeModelName string) (string, error) { - parts := strings.Split(compositeModelName, "@") - if len(parts) < 2 { - return "", fmt.Errorf("invalid model name format: %s", compositeModelName) - } - - var providerName, instanceName string - switch len(parts) { - case 2: - instanceName = "default" - providerName = parts[1] - case 3: - instanceName = parts[1] - providerName = parts[2] - default: - return "", fmt.Errorf("invalid model name format: %s", compositeModelName) - } - - provider, err := s.modelProviderDAO.GetByTenantIDAndProviderName(ctx, dao.DB, tenantID, providerName) - if err != nil { - return "", fmt.Errorf("provider %q not found: %w", providerName, err) - } - if provider == nil { - return "", fmt.Errorf("provider %q not found", providerName) - } - - instance, err := s.modelInstanceDAO.GetByProviderIDAndInstanceName(ctx, dao.DB, provider.ID, instanceName) - if err != nil { - return "", fmt.Errorf("instance %q not found: %w", instanceName, err) - } - if instance == nil { - return "", fmt.Errorf("instance %q not found", instanceName) - } - - if instance.APIKey == "" { - return "", fmt.Errorf("no API key configured for model %s", compositeModelName) - } - return instance.APIKey, nil -} - // GetTenantList get tenant list for a user func (s *TenantService) GetTenantList(ctx context.Context, userID string) ([]*TenantListItem, error) { tenants, err := s.userTenantDAO.GetTenantsByUserID(ctx, dao.DB, userID) @@ -363,424 +257,6 @@ func (s *TenantService) DeleteChunkStore(ctx context.Context, kbID string) (comm return common.CodeSuccess, nil } -type ModelItem struct { - ModelProvider *string `json:"model_provider"` - ModelInstance *string `json:"model_instance"` - ModelName *string `json:"model_name"` - ModelID string `json:"model_id"` - ModelType string `json:"model_type"` - Enable bool `json:"enable"` -} - -func tenantDefaultModelFields(modelType string) (string, string, entity.ModelType, error) { - switch modelType { - case "chat": - return "llm_id", "tenant_llm_id", entity.ModelTypeChat, nil - case "embedding": - return "embd_id", "tenant_embd_id", entity.ModelTypeEmbedding, nil - case "rerank": - return "rerank_id", "tenant_rerank_id", entity.ModelTypeRerank, nil - case "asr", "speech2text": - return "asr_id", "tenant_asr_id", entity.ModelTypeSpeech2Text, nil - case "vision", "image2text": - return "img2txt_id", "tenant_img2txt_id", entity.ModelTypeImage2Text, nil - case "tts": - return "tts_id", "tenant_tts_id", entity.ModelTypeTTS, nil - case "ocr": - return "ocr_id", "tenant_ocr_id", entity.ModelTypeOCR, nil - case "decision": - // The decision model is its own type rather than a chat model: it is - // spoken to in the System One protocol (one request, several questions) - // and is chosen independently of the answering model. - return "decision_id", "tenant_decision_id", entity.ModelTypeDecision, nil - default: - return "", "", 0, fmt.Errorf("model type %s is invalid", modelType) - } -} - -func ptrStringValue(s *string) string { - if s == nil { - return "" - } - return *s -} - -// GetDefaultModelName returns the selected model reference for a tenant. Request -// contexts use that user's membership; trusted background work uses the owner. -func (s *TenantService) GetDefaultModelName(ctx context.Context, tenantID string, modelType entity.ModelType) (string, error) { - return s.GetDefaultModelNameForTenant(ctx, tenantID, modelType) -} - -// GetDefaultModelNameForUser returns the user's setting in the current Tenant. -func (s *TenantService) GetDefaultModelNameForUser(ctx context.Context, userID string, modelType entity.ModelType) (string, error) { - membership, _, err := s.currentUserTenant(ctx, userID) - if err != nil { - return "", err - } - if !isDefaultModelType(modelType) { - return "", fmt.Errorf("invalid model type: %s", modelType) - } - modelName, _ := defaultModelRefs(membership, modelType) - return modelName, nil -} - -// GetDefaultModelNameForTenant resolves the request user's membership setting -// when present and falls back to the tenant owner for background work. -func (s *TenantService) GetDefaultModelNameForTenant(ctx context.Context, tenantID string, modelType entity.ModelType) (string, error) { - if !isDefaultModelType(modelType) { - return "", fmt.Errorf("invalid model type: %s", modelType) - } - membership, err := s.defaultModelMembershipForTenant(ctx, tenantID) - if err != nil { - return "", err - } - modelName, _ := defaultModelRefs(membership, modelType) - return modelName, nil -} - -// GetDefaultModelRefsForUser returns the model reference and resolved tenant -// model ID configured by userID in tenantID. -func (s *TenantService) GetDefaultModelRefsForUser(ctx context.Context, userID, tenantID string, modelType entity.ModelType) (string, string, error) { - if !isDefaultModelType(modelType) { - return "", "", fmt.Errorf("invalid model type: %s", modelType) - } - membership, err := s.defaultModelMembership(ctx, userID, tenantID) - if err != nil { - return "", "", err - } - modelName, modelID := defaultModelRefs(membership, modelType) - return modelName, modelID, nil -} - -// GetDefaultModelRefsForTenant resolves model settings for the request user or, -// without a request identity, the tenant owner. -func (s *TenantService) GetDefaultModelRefsForTenant(ctx context.Context, tenantID string, modelType entity.ModelType) (string, string, error) { - if !isDefaultModelType(modelType) { - return "", "", fmt.Errorf("invalid model type: %s", modelType) - } - membership, err := s.defaultModelMembershipForTenant(ctx, tenantID) - if err != nil { - return "", "", err - } - modelName, modelID := defaultModelRefs(membership, modelType) - return modelName, modelID, nil -} - -func isDefaultModelType(modelType entity.ModelType) bool { - switch modelType { - case entity.ModelTypeChat, entity.ModelTypeEmbedding, entity.ModelTypeRerank, - entity.ModelTypeSpeech2Text, entity.ModelTypeImage2Text, entity.ModelTypeTTS, - entity.ModelTypeOCR, entity.ModelTypeDecision: - return true - default: - return false - } -} - -func (s *TenantService) defaultModelMembership(ctx context.Context, userID, tenantID string) (*entity.UserTenant, error) { - if strings.TrimSpace(userID) == "" { - return s.userTenantDAO.GetOwnerByTenantID(ctx, dao.DB, tenantID) - } - membership, err := s.userTenantDAO.FilterByUserIDAndTenantID(ctx, dao.DB, userID, tenantID) - if err == nil { - return membership, nil - } - if !errors.Is(err, gorm.ErrRecordNotFound) || userID != tenantID { - return nil, err - } - return s.userTenantDAO.GetOwnerByTenantID(ctx, dao.DB, tenantID) -} - -func (s *TenantService) defaultModelMembershipForTenant(ctx context.Context, tenantID string) (*entity.UserTenant, error) { - if subject, ok := permission.SubjectFromContext(ctx); ok && strings.TrimSpace(subject.UserID) != "" && - (strings.TrimSpace(subject.TenantID) == "" || strings.TrimSpace(subject.TenantID) == tenantID) { - return s.defaultModelMembership(ctx, subject.UserID, tenantID) - } - return s.userTenantDAO.GetOwnerByTenantID(ctx, dao.DB, tenantID) -} - -func (s *TenantService) GetModelInfo(ctx context.Context, tenantID string, defaultModel string, modelType string) (*string, *string, *string, bool, error) { - // Mirror Python's _get_model_info: right-anchored rsplit so that model - // names containing '@' (e.g. LM Studio IDs like - // "text-embedding-nomic-embed-text-v1.5@q8_0") remain intact. - // The composite key is: modelName@instanceName@providerName or - // modelName@providerName. - parts := rsplitN(defaultModel, "@", 2) - var modelName, instanceName, providerName string - switch len(parts) { - case 3: - modelName, instanceName, providerName = parts[0], parts[1], parts[2] - case 2: - modelName, providerName = parts[0], parts[1] - instanceName = "default" - default: - modelName = parts[0] - providerName = "" - instanceName = "default" - } - - // Special case: OCR with infiniflow@default@deepdoc is always enabled. - if modelType == "ocr" && providerName == "infiniflow" && instanceName == "default" && modelName == "deepdoc" { - return &providerName, &instanceName, &modelName, true, nil - } - - // Special case: TEI Builtin embedding model. - composeProfiles := common.GetEnv(common.EnvComposeProfiles) - teiModel := common.GetEnv(common.EnvTEIModel) - if modelType == "embedding" && strings.Contains(composeProfiles, "tei-") && teiModel != "" && - modelName == teiModel && (providerName == "" || providerName == "Builtin") { - return &providerName, &instanceName, &modelName, true, nil - } - - // Check if the provider exists for the tenant. - modelProvider, err := s.modelProviderDAO.GetByTenantIDAndProviderName(ctx, dao.DB, tenantID, providerName) - if err != nil { - return nil, nil, nil, false, err - } - - // Check if the instance exists. - modelInstance, err := s.modelInstanceDAO.GetByProviderIDAndInstanceName(ctx, dao.DB, modelProvider.ID, instanceName) - if err != nil { - return nil, nil, nil, false, err - } - - // Check if the model exists and is active. Model type validity is - // enforced on the save path against the tenant_model record, so models - // added online that are absent from the static catalog still resolve. - modelEntity, err := s.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(ctx, dao.DB, modelProvider.ID, modelInstance.ID, modelName) - if err != nil { - if !dao.IsNotFoundErr(err) { - return nil, nil, nil, false, err - } - } - if modelEntity == nil { - return nil, nil, nil, false, fmt.Errorf("model %s isn't available", modelName) - } - if modelEntity.Status != "active" { - return nil, nil, nil, false, fmt.Errorf("model %s isn't available", modelName) - } - - return &providerName, &instanceName, &modelName, true, nil -} - -// rsplitN splits s by sep from the right, limiting to n+1 parts (mirrors -// Python's str.rsplit(sep, maxsplit)). -func rsplitN(s, sep string, n int) []string { - if n <= 0 { - return []string{s} - } - result := make([]string, 0, n+1) - remaining := s - for i := 0; i < n; i++ { - idx := strings.LastIndex(remaining, sep) - if idx < 0 { - result = append(result, remaining) - // Reverse the collected parts. - for left, right := 0, len(result)-1; left < right; left, right = left+1, right-1 { - result[left], result[right] = result[right], result[left] - } - return result - } - result = append(result, remaining[idx+len(sep):]) - remaining = remaining[:idx] - } - result = append(result, remaining) - for left, right := 0, len(result)-1; left < right; left, right = left+1, right-1 { - result[left], result[right] = result[right], result[left] - } - return result -} - -func (s *TenantService) ListTenantDefaultModels(ctx context.Context, userID string) ([]ModelItem, error) { - membership, tenantID, err := s.currentUserTenant(ctx, userID) - if err != nil { - return nil, err - } - return listDefaultModels(ctx, s, membership, tenantID), nil -} - -func listDefaultModels(ctx context.Context, s *TenantService, membership *entity.UserTenant, tenantID string) []ModelItem { - type defaultModel struct { - modelType string - modelID string - rowID *string - } - models := []defaultModel{ - {modelType: "chat", modelID: membership.LLMID, rowID: membership.TenantLLMID}, - {modelType: "embedding", modelID: membership.EmbdID, rowID: membership.TenantEmbdID}, - {modelType: "rerank", modelID: membership.RerankID, rowID: membership.TenantRerankID}, - {modelType: "asr", modelID: membership.ASRID, rowID: membership.TenantASRID}, - {modelType: "vision", modelID: membership.Img2TxtID, rowID: membership.TenantImg2TxtID}, - {modelType: "ocr", modelID: ptrStringValue(membership.OCRID), rowID: membership.TenantOCRID}, - {modelType: "tts", modelID: ptrStringValue(membership.TTSID), rowID: membership.TenantTTSID}, - {modelType: "decision", modelID: ptrStringValue(membership.DecisionID), rowID: membership.TenantDecisionID}, - } - result := make([]ModelItem, 0, len(models)) - for _, model := range models { - if strings.TrimSpace(model.modelID) == "" { - continue - } - provider, instance, name, enabled, err := s.GetModelInfo(ctx, tenantID, model.modelID, model.modelType) - if err != nil { - continue - } - result = append(result, ModelItem{ - ModelProvider: provider, - ModelInstance: instance, - ModelName: name, - ModelID: ptrStringValue(model.rowID), - ModelType: model.modelType, - Enable: enabled, - }) - } - return result -} - -func (s *TenantService) checkModelAvailable(ctx context.Context, tenantID, providerName, instanceName, modelName, modelType string) error { - _, _, modelTypeBit, err := tenantDefaultModelFields(modelType) - if err != nil { - return err - } - - // Static bypass: deepdoc is a built-in model that doesn't need DB checks (mirrors Python _check_model_available). - if providerName == "infiniflow" && instanceName == "default" && modelName == "deepdoc" { - return nil - } - - // Static bypass: OCR with infiniflow@default@deepdoc is always enabled (mirrors Python _check_model_available). - if modelType == "ocr" && providerName == "infiniflow" && instanceName == "default" { - return nil - } - - // Static bypass: TEI Builtin embedding model when COMPOSE_PROFILES includes tei- (mirrors Python _check_model_available). - composeProfiles := common.GetEnv(common.EnvComposeProfiles) - teiModel := common.GetEnv(common.EnvTEIModel) - if modelType == "embedding" && strings.Contains(composeProfiles, "tei-") && teiModel != "" && - modelName == teiModel && (providerName == "" || providerName == "Builtin") { - return nil - } - - // Check if the provider and instance exists - modelProvider, err := s.modelProviderDAO.GetByTenantIDAndProviderName(ctx, dao.DB, tenantID, providerName) - if err != nil { - return err - } - - modelInstance, err := s.modelInstanceDAO.GetByProviderIDAndInstanceName(ctx, dao.DB, modelProvider.ID, instanceName) - if err != nil { - return err - } - - // Validate model availability through the DB (TenantModel table) - modelEntity, err := s.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(ctx, dao.DB, modelProvider.ID, modelInstance.ID, modelName) - if err != nil { - if dao.IsNotFoundErr(err) { - return fmt.Errorf("model %s isn't available", modelName) - } - return err - } - if modelEntity.Status != "active" { - return fmt.Errorf("model %s isn't available", modelName) - } - if !entity.ModelType(modelEntity.ModelType).Has(modelTypeBit) { - return fmt.Errorf("model %s isn't a %s model", modelName, modelType) - } - - return nil -} - -func (s *TenantService) SetTenantDefaultModels(ctx context.Context, userID, modelProvider, modelInstance, modelName, modelType, modelID string) error { - membership, tenantID, err := s.currentUserTenant(ctx, userID) - if err != nil { - return err - } - var defaultModel string - modelTypeID, tenantModelTypeID, modelTypeBit, err := tenantDefaultModelFields(modelType) - if err != nil { - return err - } - - var tenantModelID interface{} - if modelID != "" { - modelEntity, err := s.modelDAO.GetByID(ctx, dao.DB, modelID) - if err != nil { - return fmt.Errorf("model ID %s is invalid", modelID) - } - instanceEntity, err := s.modelInstanceDAO.GetByID(ctx, dao.DB, modelEntity.InstanceID) - if err != nil { - return fmt.Errorf("instance for model %s not found: %w", modelID, err) - } - providerEntity, err := s.modelProviderDAO.GetByID(ctx, dao.DB, instanceEntity.ProviderID) - if err != nil { - return fmt.Errorf("provider for model %s not found: %w", modelID, err) - } - if providerEntity.TenantID != tenantID { - return fmt.Errorf("model %s does not belong to your tenant", modelID) - } - if modelEntity.Status != "active" { - return fmt.Errorf("model %s isn't available", modelEntity.ModelName) - } - if !entity.ModelType(modelEntity.ModelType).Has(modelTypeBit) { - return fmt.Errorf("model %s isn't a %s model", modelEntity.ModelName, modelType) - } - modelProvider = providerEntity.ProviderName - modelInstance = instanceEntity.InstanceName - modelName = modelEntity.ModelName - tenantModelID = modelID - } - - if modelProvider == "" && modelInstance == "" && modelName == "" { - defaultModel = "" - tenantModelID = nil - } else if modelProvider != "" && modelInstance != "" && modelName != "" { - if err := s.checkModelAvailable(ctx, tenantID, modelProvider, modelInstance, modelName, modelType); err != nil { - return err - } - if modelID == "" { - if modelProvider == "Builtin" { - tenantModelID = nil - } else { - provider, err := s.modelProviderDAO.GetByTenantIDAndProviderName(ctx, dao.DB, tenantID, modelProvider) - if err != nil { - return err - } - instance, err := s.modelInstanceDAO.GetByProviderIDAndInstanceName(ctx, dao.DB, provider.ID, modelInstance) - if err != nil { - return err - } - model, err := s.modelDAO.GetModelByProviderIDAndInstanceIDAndModelName(ctx, dao.DB, provider.ID, instance.ID, modelName) - if err != nil { - return err - } - tenantModelID = model.ID - } - } - defaultModel = fmt.Sprintf("%s@%s@%s", modelName, modelInstance, modelProvider) - } else { - return fmt.Errorf("model provider, instance and name must be specified together") - } - - return dao.DB.WithContext(ctx).Model(&entity.UserTenant{}).Where("id = ?", membership.ID).Updates(map[string]interface{}{ - modelTypeID: defaultModel, - tenantModelTypeID: tenantModelID, - }).Error -} - -func (s *TenantService) currentUserTenant(ctx context.Context, userID string) (*entity.UserTenant, string, error) { - membership, handled, err := resolveCurrentUserTenantEE(ctx, userID) - if err != nil { - return nil, "", err - } - if handled { - return membership, membership.TenantID, nil - } - membership, err = s.userTenantDAO.GetOwnerByUserID(ctx, dao.DB, userID) - if err != nil { - return nil, "", err - } - return membership, membership.TenantID, nil -} - // Tenant member role constants. const ( TenantRoleOwner = "owner" diff --git a/internal/service/tenant_test.go b/internal/service/tenant_test.go index 60354633e2..2940f59aef 100644 --- a/internal/service/tenant_test.go +++ b/internal/service/tenant_test.go @@ -201,9 +201,9 @@ func TestTenantRoleConstants(t *testing.T) { } } -// TestSetTenantDefaultModels_WithModelID verifies that SetTenantDefaultModels -// correctly resolves a modelID to composite name, validates ownership, and updates the tenant. -func TestSetTenantDefaultModels_WithModelID(t *testing.T) { +// TestSetModelSettings_WithModelID verifies that model IDs resolve to canonical references +// and are persisted on the caller's active tenant membership. +func TestSetModelSettings_WithModelID(t *testing.T) { // 1. Setup SQLite in-memory DB db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ TranslateError: true, @@ -263,22 +263,13 @@ func TestSetTenantDefaultModels_WithModelID(t *testing.T) { err = db.Create(&entity.TenantModelProvider{ ID: providerID, TenantID: tenantID, + UserID: userID, ProviderName: "OpenAI", }).Error if err != nil { t.Fatalf("failed to create provider: %v", err) } - // Insert Real Instance (for checkModelAvailable lookup) - err = db.Create(&entity.TenantModelInstance{ - ID: "instance-real", - ProviderID: providerID, - InstanceName: "default", - }).Error - if err != nil { - t.Fatalf("failed to create real instance: %v", err) - } - // Insert Dummy Instance (associated with the model record) err = db.Create(&entity.TenantModelInstance{ ID: "instance-dummy", @@ -303,13 +294,14 @@ func TestSetTenantDefaultModels_WithModelID(t *testing.T) { t.Fatalf("failed to create model: %v", err) } - // 4. Run SetTenantDefaultModels - s := NewTenantService() - // Set chat model using modelID, explicitly passing "default" as instance name to bypass pre-existing checkModelAvailable panic + // 4. Run the unified Tenant model settings path. + s := NewUserTenantService() ctx := serviceTestContext(t, userID, tenantID) - err = s.SetTenantDefaultModels(ctx, userID, "", "default", "", "chat", modelID) - if err != nil { - t.Fatalf("SetTenantDefaultModels failed: %v", err) + code, err := s.SetModelSettings(ctx, userID, &SetModelSettingsRequest{ + DefaultModel: &DefaultModelSelection{ModelType: "chat", ModelID: modelID}, + }) + if err != nil || code != common.CodeSuccess { + t.Fatalf("SetModelSettings() = (%v, %v), want success", code, err) } // Verify the default is stored on the user's owner membership. diff --git a/internal/service/user.go b/internal/service/user.go index 1d784416c7..b91103e2b2 100644 --- a/internal/service/user.go +++ b/internal/service/user.go @@ -743,138 +743,6 @@ func (s *UserService) ChangePassword(ctx context.Context, user *entity.User, req return common.CodeSuccess, nil } -// SetTenantInfoRequest represents the request for setting tenant info -type SetTenantInfoRequest struct { - TenantID *string `json:"tenant_id"` - ASRID *string `json:"asr_id"` - EmbdID *string `json:"embd_id"` - Img2TxtID *string `json:"img2txt_id"` - LLMID *string `json:"llm_id"` - RerankID *string `json:"rerank_id"` - TTSID *string `json:"tts_id"` - Raw map[string]interface{} `json:"-"` -} - -// SetTenantInfo updates default models on the caller's current Tenant membership. -func (s *UserService) SetTenantInfo(ctx context.Context, userID string, req *SetTenantInfoRequest) (common.ErrorCode, error) { - if req == nil { - return common.CodeArgumentError, fmt.Errorf("model settings are required") - } - updates, err := tenantModelSettingUpdates(req.Raw) - if err != nil { - return common.CodeArgumentError, err - } - tenantID, err := requireCurrentTenantID(ctx, userID) - if err != nil { - return common.CodeForbidden, err - } - if req.TenantID != nil && strings.TrimSpace(*req.TenantID) != "" && strings.TrimSpace(*req.TenantID) != tenantID { - return common.CodeForbidden, fmt.Errorf("tenant_id must match the active Tenant context") - } - userTenantDAO := dao.NewUserTenantDAO() - membership, err := userTenantDAO.FilterByUserIDAndTenantID(ctx, dao.DB, userID, tenantID) - if err != nil { - return common.CodeForbidden, err - } - if len(updates) > 0 { - if err := dao.DB.WithContext(ctx).Model(&entity.UserTenant{}).Where("id = ?", membership.ID).Updates(updates).Error; err != nil { - return common.CodeExceptionError, err - } - } - return common.CodeSuccess, nil -} - -func tenantModelSettingUpdates(raw map[string]interface{}) (map[string]interface{}, error) { - updates := make(map[string]interface{}, len(raw)) - for key, value := range raw { - if key == "tenant_id" { - if _, ok := value.(string); !ok { - return nil, fmt.Errorf("tenant_id must be a string") - } - continue - } - nullable := false - switch key { - case "llm_id", "embd_id", "asr_id", "img2txt_id", "rerank_id": - case "tenant_llm_id", "tenant_embd_id", "tenant_asr_id", "tenant_img2txt_id", "tenant_rerank_id", - "tts_id", "tenant_tts_id", "ocr_id", "tenant_ocr_id", "decision_id", "tenant_decision_id": - nullable = true - default: - return nil, fmt.Errorf("unsupported model setting %q", key) - } - if value == nil { - if !nullable { - return nil, fmt.Errorf("%s cannot be null", key) - } - updates[key] = nil - continue - } - if _, ok := value.(string); !ok { - return nil, fmt.Errorf("%s must be a string", key) - } - updates[key] = value - } - return updates, nil -} - -// UserTenantService user tenant service -// Provides business logic for user-tenant relationship management -type UserTenantService struct { - userTenantDAO *dao.UserTenantDAO -} - -// NewUserTenantService creates a new UserTenantService instance -/** - * Returns: - * - *UserTenantService: a new UserTenantService instance - */ -func NewUserTenantService() *UserTenantService { - return &UserTenantService{ - userTenantDAO: dao.NewUserTenantDAO(), - } -} - -// UserTenantRelation represents a user-tenant relationship response -// This structure matches the Python implementation's return format -type UserTenantRelation struct { - ID string `json:"id"` - UserID string `json:"user_id"` - TenantID string `json:"tenant_id"` - Role string `json:"role"` -} - -// GetUserTenantRelationByUserIDWithContext retrieves all user-tenant relationships for a given user ID with context. -func (s *UserTenantService) GetUserTenantRelationByUserIDWithContext(ctx context.Context, userID string) ([]*UserTenantRelation, error) { - relations, err := s.userTenantDAO.GetByUserID(ctx, dao.DB, userID) - if err != nil { - return nil, err - } - - result := make([]*UserTenantRelation, len(relations)) - for i, rel := range relations { - result[i] = convertToUserTenantRelation(rel) - } - - return result, nil -} - -// convertToUserTenantRelation converts model.UserTenant to UserTenantRelation -/** - * Parameters: - * - userTenant: the model.UserTenant to convert - * - * Returns: - * - *UserTenantRelation: the converted UserTenantRelation - */ -func convertToUserTenantRelation(userTenant *entity.UserTenant) *UserTenantRelation { - return &UserTenantRelation{ - ID: userTenant.ID, - UserID: userTenant.UserID, - TenantID: userTenant.TenantID, - Role: userTenant.Role, - } -} - // ResolveAPIToken validates an API key and resolves its authentication identity. func (s *UserService) ResolveAPIToken(ctx context.Context, authorization string) (*AuthIdentity, common.ErrorCode, error) { token, err := parseAuthorizationToken(authorization) diff --git a/internal/service/user_set_tenant_info_test.go b/internal/service/user_set_tenant_info_test.go index 688c0ecfb5..d1c5f32b17 100644 --- a/internal/service/user_set_tenant_info_test.go +++ b/internal/service/user_set_tenant_info_test.go @@ -22,7 +22,7 @@ import ( "ragflow/internal/common" ) -func TestSetTenantInfoRejectsInvalidSettingTypesBeforeDatabaseLookup(t *testing.T) { +func TestSetModelSettingsRejectsInvalidSettingTypesBeforeDatabaseLookup(t *testing.T) { tests := []struct { name string raw map[string]interface{} @@ -36,16 +36,16 @@ func TestSetTenantInfoRejectsInvalidSettingTypesBeforeDatabaseLookup(t *testing. for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - code, err := (&UserService{}).SetTenantInfo(context.Background(), "user-1", &SetTenantInfoRequest{Raw: tt.raw}) + code, err := (&UserTenantService{}).SetModelSettings(context.Background(), "user-1", &SetModelSettingsRequest{Raw: tt.raw}) if code != common.CodeArgumentError || err == nil { - t.Fatalf("SetTenantInfo() = (%v, %v), want argument error", code, err) + t.Fatalf("SetModelSettings() = (%v, %v), want argument error", code, err) } }) } } -func TestTenantModelSettingUpdatesAllowsNullOnlyForNullableColumns(t *testing.T) { - updates, err := tenantModelSettingUpdates(map[string]interface{}{ +func TestModelSettingUpdatesAllowsNullOnlyForNullableColumns(t *testing.T) { + updates, err := modelSettingUpdates(map[string]interface{}{ "llm_id": "chat@provider", "tenant_llm_id": nil, "tts_id": nil, @@ -56,7 +56,7 @@ func TestTenantModelSettingUpdatesAllowsNullOnlyForNullableColumns(t *testing.T) "tenant_decision_id": nil, }) if err != nil { - t.Fatalf("tenantModelSettingUpdates() error = %v", err) + t.Fatalf("modelSettingUpdates() error = %v", err) } if len(updates) != 8 { t.Fatalf("len(updates) = %d, want 8", len(updates)) diff --git a/internal/service/user_tenant.go b/internal/service/user_tenant.go new file mode 100644 index 0000000000..e218bce42e --- /dev/null +++ b/internal/service/user_tenant.go @@ -0,0 +1,389 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "context" + "errors" + "fmt" + "strings" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/entity" + "ragflow/internal/permission" +) + +// SetModelSettingsRequest represents default model selections for a user in the current tenant. +type SetModelSettingsRequest struct { + TenantID *string `json:"tenant_id"` + ASRID *string `json:"asr_id"` + EmbdID *string `json:"embd_id"` + Img2TxtID *string `json:"img2txt_id"` + LLMID *string `json:"llm_id"` + RerankID *string `json:"rerank_id"` + TTSID *string `json:"tts_id"` + Raw map[string]interface{} `json:"-"` + DefaultModel *DefaultModelSelection `json:"-"` +} + +// DefaultModelSelection is the single-model form used by the default-model endpoints. +type DefaultModelSelection struct { + ModelProvider string + ModelInstance string + ModelName string + ModelType string + ModelID string +} + +// UserTenantService manages user settings scoped to a tenant membership. +type UserTenantService struct { + userTenantDAO *dao.UserTenantDAO + tenantDAO *dao.TenantDAO +} + +// NewUserTenantService creates a UserTenantService. +func NewUserTenantService() *UserTenantService { + return &UserTenantService{ + userTenantDAO: dao.NewUserTenantDAO(), + tenantDAO: dao.NewTenantDAO(), + } +} + +// TenantInfoResponse contains the current tenant and membership settings. +type TenantInfoResponse struct { + TenantID string `json:"tenant_id"` + Name *string `json:"name,omitempty"` + LLMID string `json:"llm_id"` + EmbDID string `json:"embd_id"` + RerankID string `json:"rerank_id"` + ASRID string `json:"asr_id"` + Img2TxtID string `json:"img2txt_id"` + TTSID string `json:"tts_id"` + ParserIDs string `json:"parser_ids"` + Role string `json:"role"` +} + +// GetCurrentTenantInfo returns the active tenant with the user's membership settings. +func (s *UserTenantService) GetCurrentTenantInfo(ctx context.Context, userID string) (*TenantInfoResponse, error) { + tenantID, err := requireCurrentTenantID(ctx, userID) + if err != nil { + return nil, err + } + membership, err := s.userTenantDAO.FilterByUserIDAndTenantID(ctx, dao.DB, userID, tenantID) + if err != nil { + return nil, err + } + tenant, err := s.tenantDAO.GetByID(ctx, dao.DB, tenantID) + if err != nil { + return nil, err + } + return &TenantInfoResponse{ + TenantID: tenant.ID, + Name: tenant.Name, + LLMID: membership.LLMID, + EmbDID: membership.EmbdID, + RerankID: membership.RerankID, + ASRID: membership.ASRID, + Img2TxtID: membership.Img2TxtID, + TTSID: ptrStringValue(membership.TTSID), + ParserIDs: tenant.ParserIDs, + Role: membership.Role, + }, nil +} + +// UserTenantRelation represents a user-tenant relationship response. +type UserTenantRelation struct { + ID string `json:"id"` + UserID string `json:"user_id"` + TenantID string `json:"tenant_id"` + Role string `json:"role"` +} + +// GetUserTenantRelationByUserIDWithContext retrieves a user's active tenant relationships. +func (s *UserTenantService) GetUserTenantRelationByUserIDWithContext(ctx context.Context, userID string) ([]*UserTenantRelation, error) { + relations, err := s.userTenantDAO.GetByUserID(ctx, dao.DB, userID) + if err != nil { + return nil, err + } + + result := make([]*UserTenantRelation, len(relations)) + for i, relation := range relations { + result[i] = &UserTenantRelation{ + ID: relation.ID, + UserID: relation.UserID, + TenantID: relation.TenantID, + Role: relation.Role, + } + } + return result, nil +} + +var defaultModelTypeNames = [...]string{"chat", "embedding", "rerank", "asr", "vision", "tts", "ocr", "decision"} + +// SetModelSettings validates and updates default models on the user's active tenant membership. +func (s *UserTenantService) SetModelSettings(ctx context.Context, userID string, req *SetModelSettingsRequest) (common.ErrorCode, error) { + if req == nil || (req.Raw != nil && req.DefaultModel != nil) { + return common.CodeArgumentError, fmt.Errorf("provide either model settings or one default model selection") + } + raw := req.Raw + if selection := req.DefaultModel; selection != nil { + legacyField, idField, _, err := defaultModelFields(selection.ModelType) + if err != nil { + return common.CodeArgumentError, err + } + provider := strings.TrimSpace(selection.ModelProvider) + instance := strings.TrimSpace(selection.ModelInstance) + name := strings.TrimSpace(selection.ModelName) + modelID := strings.TrimSpace(selection.ModelID) + anyNameField := provider != "" || instance != "" || name != "" + if anyNameField && modelID == "" && (provider == "" || instance == "" || name == "") { + return common.CodeArgumentError, fmt.Errorf("model provider, instance and name must be specified together") + } + modelRef := "" + if anyNameField && provider != "" && instance != "" && name != "" { + modelRef = fmt.Sprintf("%s@%s@%s", name, instance, provider) + } + var modelIDValue interface{} + if modelID != "" { + modelIDValue = modelID + } + raw = map[string]interface{}{legacyField: modelRef, idField: modelIDValue} + } + updates, err := modelSettingUpdates(raw) + if err != nil { + return common.CodeArgumentError, err + } + tenantID, err := requireCurrentTenantID(ctx, userID) + if err != nil { + return common.CodeForbidden, err + } + if req.TenantID != nil && strings.TrimSpace(*req.TenantID) != "" && strings.TrimSpace(*req.TenantID) != tenantID { + return common.CodeForbidden, fmt.Errorf("tenant_id must match the active Tenant context") + } + if rawTenantID, ok := raw["tenant_id"].(string); ok && strings.TrimSpace(rawTenantID) != "" && strings.TrimSpace(rawTenantID) != tenantID { + return common.CodeForbidden, fmt.Errorf("tenant_id must match the active Tenant context") + } + if s == nil || s.userTenantDAO == nil || dao.DB == nil { + return common.CodeExceptionError, fmt.Errorf("User Tenant service is not initialized") + } + membership, err := s.userTenantDAO.FilterByUserIDAndTenantID(ctx, dao.DB, userID, tenantID) + if err != nil { + if dao.IsNotFoundErr(err) { + return common.CodeForbidden, err + } + return common.CodeExceptionError, err + } + + factory := NewModelFactory() + modelUpdates := make(map[string]interface{}, len(updates)) + for _, modelTypeName := range defaultModelTypeNames { + legacyField, idField, modelType, _ := defaultModelFields(modelTypeName) + legacyValue, hasLegacy := updates[legacyField] + idValue, hasID := updates[idField] + if !hasLegacy && !hasID { + continue + } + legacyRef, _ := legacyValue.(string) + idRef, _ := idValue.(string) + legacyRef = strings.TrimSpace(legacyRef) + idRef = strings.TrimSpace(idRef) + modelRef := idRef + if modelRef == "" { + modelRef = legacyRef + } + if modelRef == "" { + modelUpdates[legacyField] = "" + modelUpdates[idField] = nil + continue + } + + info, err := factory.ResolveInfo(ctx, Identity{UserID: userID, TenantID: tenantID}, modelType, modelRef) + if err != nil { + switch { + case errors.Is(err, permission.ErrPermissionDenied), errors.Is(err, permission.ErrResourceNotFound), errors.Is(err, permission.ErrMembershipNotFound): + return common.CodeForbidden, err + case errors.Is(err, errModelConfigUnavailable): + return common.CodeArgumentError, err + default: + return common.CodeExceptionError, err + } + } + if info == nil { + return common.CodeArgumentError, fmt.Errorf("model %q is unavailable", modelRef) + } + if info.InstanceName != "" { + modelUpdates[legacyField] = fmt.Sprintf("%s@%s@%s", info.Name, info.InstanceName, info.ProviderName) + } else if info.ProviderName != "" { + modelUpdates[legacyField] = fmt.Sprintf("%s@%s", info.Name, info.ProviderName) + } else { + modelUpdates[legacyField] = info.Name + } + var resolvedModelID interface{} + if info.ID != "" { + resolvedModelID = info.ID + } + modelUpdates[idField] = resolvedModelID + } + if len(modelUpdates) == 0 { + return common.CodeSuccess, nil + } + if err := dao.DB.WithContext(ctx).Model(&entity.UserTenant{}).Where("id = ?", membership.ID).Updates(modelUpdates).Error; err != nil { + return common.CodeExceptionError, err + } + return common.CodeSuccess, nil +} + +func modelSettingField(key string) (string, string, bool) { + for _, modelType := range defaultModelTypeNames { + legacyField, idField, _, err := defaultModelFields(modelType) + if err != nil { + continue + } + if key == legacyField || key == idField { + nullable := key == idField || legacyField == "tts_id" || legacyField == "ocr_id" || legacyField == "decision_id" + return legacyField, idField, nullable + } + } + return "", "", false +} + +func modelSettingUpdates(raw map[string]interface{}) (map[string]interface{}, error) { + updates := make(map[string]interface{}, len(raw)) + for key, value := range raw { + if key == "tenant_id" { + if _, ok := value.(string); !ok { + return nil, fmt.Errorf("tenant_id must be a string") + } + continue + } + legacyField, _, nullable := modelSettingField(key) + if legacyField == "" { + return nil, fmt.Errorf("unsupported model setting %q", key) + } + if value == nil { + if !nullable { + return nil, fmt.Errorf("%s cannot be null", key) + } + updates[key] = nil + continue + } + if _, ok := value.(string); !ok { + return nil, fmt.Errorf("%s must be a string", key) + } + updates[key] = value + } + return updates, nil +} + +// ModelItem describes one configured default model. +type ModelItem struct { + ModelProvider *string `json:"model_provider"` + ModelInstance *string `json:"model_instance"` + ModelName *string `json:"model_name"` + ModelID string `json:"model_id"` + ModelType string `json:"model_type"` + Enable bool `json:"enable"` +} + +func defaultModelFields(modelType string) (string, string, entity.ModelType, error) { + switch modelType { + case "chat": + return "llm_id", "tenant_llm_id", entity.ModelTypeChat, nil + case "embedding": + return "embd_id", "tenant_embd_id", entity.ModelTypeEmbedding, nil + case "rerank": + return "rerank_id", "tenant_rerank_id", entity.ModelTypeRerank, nil + case "asr", "speech2text": + return "asr_id", "tenant_asr_id", entity.ModelTypeSpeech2Text, nil + case "vision", "image2text": + return "img2txt_id", "tenant_img2txt_id", entity.ModelTypeImage2Text, nil + case "tts": + return "tts_id", "tenant_tts_id", entity.ModelTypeTTS, nil + case "ocr": + return "ocr_id", "tenant_ocr_id", entity.ModelTypeOCR, nil + case "decision": + return "decision_id", "tenant_decision_id", entity.ModelTypeDecision, nil + default: + return "", "", 0, fmt.Errorf("model type %s is invalid", modelType) + } +} + +func ptrStringValue(value *string) string { + if value == nil { + return "" + } + return *value +} + +// ListDefaultModels returns the models configured on the user's current tenant membership. +func (s *UserTenantService) ListDefaultModels(ctx context.Context, userID string) ([]ModelItem, error) { + tenantID, err := requireCurrentTenantID(ctx, userID) + if err != nil { + return nil, err + } + membership, err := s.userTenantDAO.FilterByUserIDAndTenantID(ctx, dao.DB, userID, tenantID) + if err != nil { + return nil, err + } + type defaultModel struct { + modelType string + modelName string + modelID *string + } + models := []defaultModel{ + {modelType: "chat", modelName: membership.LLMID, modelID: membership.TenantLLMID}, + {modelType: "embedding", modelName: membership.EmbdID, modelID: membership.TenantEmbdID}, + {modelType: "rerank", modelName: membership.RerankID, modelID: membership.TenantRerankID}, + {modelType: "asr", modelName: membership.ASRID, modelID: membership.TenantASRID}, + {modelType: "vision", modelName: membership.Img2TxtID, modelID: membership.TenantImg2TxtID}, + {modelType: "ocr", modelName: ptrStringValue(membership.OCRID), modelID: membership.TenantOCRID}, + {modelType: "tts", modelName: ptrStringValue(membership.TTSID), modelID: membership.TenantTTSID}, + {modelType: "decision", modelName: ptrStringValue(membership.DecisionID), modelID: membership.TenantDecisionID}, + } + factory := NewModelFactory() + result := make([]ModelItem, 0, len(models)) + for _, model := range models { + modelRef := ptrStringValue(model.modelID) + if modelRef == "" { + modelRef = model.modelName + } + if strings.TrimSpace(modelRef) == "" { + continue + } + _, _, modelType, err := defaultModelFields(model.modelType) + if err != nil { + continue + } + info, err := factory.ResolveInfo(ctx, Identity{UserID: userID, TenantID: tenantID}, modelType, modelRef) + if err != nil || info == nil { + continue + } + modelID := info.ID + if modelID == "" { + modelID = ptrStringValue(model.modelID) + } + result = append(result, ModelItem{ + ModelProvider: &info.ProviderName, + ModelInstance: &info.InstanceName, + ModelName: &info.Name, + ModelID: modelID, + ModelType: model.modelType, + Enable: true, + }) + } + return result, nil +}