zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-06-29 15:31:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat[Go]: add modelID for delete_model and update_status (#16025)

Browse Source 
### What problem does this PR solve?

1. add modelID for delete_model and update_status
2. fix the bug when update-status delete model

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Haruko386
2026-06-18 17:56:51 +08:00
committed by GitHub
parent 3eb49ca7f8
commit 6beae949d8
7 changed files with 651 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
internal/handler/providers.go
View File

@@ -704,7 +704,8 @@ func (h *ProviderHandler) ListInstanceModels(c *gin.Context) {
}
type EnableOrDisableModelRequest struct {
Status string `json:"status" binding:"required"`
ModelID string `json:"model_id"`
Status string `json:"status"`
}
func (h *ProviderHandler) EnableOrDisableModel(c *gin.Context) {
@@ -726,18 +727,6 @@ func (h *ProviderHandler) EnableOrDisableModel(c *gin.Context) {
return
}
modelName := c.Param("model_name")
if modelName != "" {
modelName = strings.TrimPrefix(modelName, "/")
}
if modelName == "" {
c.JSON(http.StatusBadRequest, gin.H{
"code": 400,
"message": "Model name is required",
})
return
}
var req EnableOrDisableModelRequest
if err := c.ShouldBindJSON(&req); err != nil {
println("JSON bind error: %v (type: %T)", err, err)
@@ -749,11 +738,30 @@ func (h *ProviderHandler) EnableOrDisableModel(c *gin.Context) {
}
userID := c.GetString("user_id")
modelID := strings.TrimSpace(req.ModelID)
modelName := strings.TrimPrefix(c.Param("model_name"), "/")
modelName = strings.TrimSpace(modelName)
if modelName == "" && modelID == "" {
c.JSON(http.StatusBadRequest, gin.H{
"code": common.CodeBadRequest,
"message": "model_name or model_id is required",
})
return
}
_, err := h.modelProviderService.UpdateModelStatus(providerName, instanceName, modelName, userID, req.Status)
status := strings.TrimSpace(req.Status)
if status != "active" && status != "inactive" {
c.JSON(http.StatusBadRequest, gin.H{
"code": common.CodeBadRequest,
"message": "Status must be active or inactive",
})
return
}
code, err := h.modelProviderService.UpdateModelStatus(providerName, instanceName, modelName, userID, modelID, status)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"code": common.CodeServerError,
"code": code,
"message": err.Error(),
})
return
@@ -846,7 +854,8 @@ func (h *ProviderHandler) AddModel(c *gin.Context) {
}
type DropInstanceModelRequest struct {
Models []string `json:"models" binding:"required"`
ModelIDs []string `json:"model_ids"`
Models []string `json:"models"`
}
func (h *ProviderHandler) DropInstanceModels(c *gin.Context) {
@@ -875,13 +884,20 @@ func (h *ProviderHandler) DropInstanceModels(c *gin.Context) {
})
return
}
if len(req.ModelIDs) == 0 && len(req.Models) == 0 {
c.JSON(http.StatusBadRequest, gin.H{
"code": common.CodeBadRequest,
"message": "model_ids or models is required",
})
return
}
userID := c.GetString("user_id")
_, err := h.modelProviderService.DropInstanceModels(providerName, instanceName, userID, req.Models)
code, err := h.modelProviderService.DropInstanceModels(providerName, instanceName, userID, req.ModelIDs, req.Models)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"code": common.CodeServerError,
"code": code,
"message": err.Error(),
})
return

155
internal/handler/providers_test.go
View File

@@ -1 +1,156 @@
package handler
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"gorm.io/gorm"
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/entity"
"ragflow/internal/service"
)
func setupProviderHandlerTestDB(t *testing.T) *gorm.DB {
t.Helper()
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true})
if err != nil {
t.Fatalf("failed to open sqlite: %v", err)
}
if err := db.AutoMigrate(
&entity.UserTenant{},
&entity.TenantModelProvider{},
&entity.TenantModelInstance{},
&entity.TenantModel{},
); err != nil {
t.Fatalf("failed to migrate provider handler tables: %v", err)
}
return db
}
func useProviderHandlerTestDB(t *testing.T, db *gorm.DB) {
t.Helper()
orig := dao.DB
dao.DB = db
t.Cleanup(func() { dao.DB = orig })
}
func seedProviderHandlerModel(t *testing.T, db *gorm.DB) {
t.Helper()
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-1", TenantID: "tenant-1", ProviderName: "OpenAI"},
&entity.TenantModelInstance{ID: "instance-1", ProviderID: "provider-1", InstanceName: "default", APIKey: "sk-test", Status: "active", Extra: "{}"},
&entity.TenantModel{ID: "model-1", ProviderID: "provider-1", InstanceID: "instance-1", ModelName: "gpt-test", ModelType: "chat", Status: "active"},
}
for _, row := range rows {
if err := db.Create(row).Error; err != nil {
t.Fatalf("failed to seed %T: %v", row, err)
}
}
}
func newProviderHandlerRequest(t *testing.T, body map[string]interface{}, params ...gin.Param) (*gin.Context, *httptest.ResponseRecorder) {
t.Helper()
gin.SetMode(gin.TestMode)
payload, err := json.Marshal(body)
if err != nil {
t.Fatalf("marshal body: %v", err)
}
req := httptest.NewRequest(http.MethodPatch, "/providers/OpenAI/instances/default/models/gpt-test", bytes.NewReader(payload))
req.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = req
ctx.Params = params
ctx.Set("user_id", "user-1")
return ctx, recorder
}
func decodeProviderHandlerResponse(t *testing.T, recorder *httptest.ResponseRecorder) map[string]interface{} {
t.Helper()
var body map[string]interface{}
if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil {
t.Fatalf("decode response %q: %v", recorder.Body.String(), err)
}
return body
}
func TestProviderHandlerEnableOrDisableModelRejectsMissingModelSelector(t *testing.T) {
ctx, recorder := newProviderHandlerRequest(
t,
map[string]interface{}{"status": "active"},
gin.Param{Key: "provider_name", Value: "OpenAI"},
gin.Param{Key: "instance_name", Value: "default"},
)
NewProviderHandler(nil, service.NewModelProviderService()).EnableOrDisableModel(ctx)
if recorder.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusBadRequest, recorder.Body.String())
}
body := decodeProviderHandlerResponse(t, recorder)
if common.ErrorCode(body["code"].(float64)) != common.CodeBadRequest {
t.Fatalf("code = %v, want %v", body["code"], common.CodeBadRequest)
}
}
func TestProviderHandlerEnableOrDisableModelRejectsInvalidStatus(t *testing.T) {
ctx, recorder := newProviderHandlerRequest(
t,
map[string]interface{}{"status": "disabled"},
gin.Param{Key: "provider_name", Value: "OpenAI"},
gin.Param{Key: "instance_name", Value: "default"},
gin.Param{Key: "model_name", Value: "gpt-test"},
)
NewProviderHandler(nil, service.NewModelProviderService()).EnableOrDisableModel(ctx)
if recorder.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusBadRequest, recorder.Body.String())
}
body := decodeProviderHandlerResponse(t, recorder)
if common.ErrorCode(body["code"].(float64)) != common.CodeBadRequest {
t.Fatalf("code = %v, want %v", body["code"], common.CodeBadRequest)
}
}
func TestProviderHandlerEnableOrDisableModelUpdatesStatus(t *testing.T) {
db := setupProviderHandlerTestDB(t)
useProviderHandlerTestDB(t, db)
seedProviderHandlerModel(t, db)
ctx, recorder := newProviderHandlerRequest(
t,
map[string]interface{}{"status": "inactive"},
gin.Param{Key: "provider_name", Value: "OpenAI"},
gin.Param{Key: "instance_name", Value: "default"},
gin.Param{Key: "model_name", Value: "gpt-test"},
)
NewProviderHandler(nil, service.NewModelProviderService()).EnableOrDisableModel(ctx)
if recorder.Code != http.StatusOK {
t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
}
body := decodeProviderHandlerResponse(t, recorder)
if common.ErrorCode(body["code"].(float64)) != common.CodeSuccess {
t.Fatalf("code = %v, want %v", body["code"], common.CodeSuccess)
}
var got entity.TenantModel
if err := db.Where("id = ?", "model-1").First(&got).Error; err != nil {
t.Fatalf("failed to reload model: %v", err)
}
if got.Status != "inactive" {
t.Fatalf("status = %q, want inactive", got.Status)
}
}