diff --git a/admin/client/ragflow_client.py b/admin/client/ragflow_client.py index b1d1e054aa..8b67537b1a 100644 --- a/admin/client/ragflow_client.py +++ b/admin/client/ragflow_client.py @@ -691,9 +691,9 @@ class RAGFlowClient: if res_json["code"] == 0: print(res_json["data"]) else: - print(f"Fail to show license, code: {res_json['code']}, message: {res_json['message']}") + print(f"Invalid, code: {res_json['code']}, message: {res_json['data']}") else: - print(f"Fail to show license, code: {response.status_code}, body: {response.text}") + print(f"Fail to check license, code: {response.status_code}, body: {response.text}") def list_server_configs(self, command): """List server configs by calling /system/configs API and flattening the JSON response.""" diff --git a/api/apps/restful_apis/chat_api.py b/api/apps/restful_apis/chat_api.py index 24f5c0259f..cf04953602 100644 --- a/api/apps/restful_apis/chat_api.py +++ b/api/apps/restful_apis/chat_api.py @@ -141,15 +141,15 @@ def _has_knowledge_placeholder(prompt_config): def _validate_name(name, *, required=True): if name is None: if required: - return None, "`name` is required." + return None, "`name` is required" return None, None if not isinstance(name, str): - return None, "Chat name must be a string." + return None, "chat name must be a string" name = name.strip() if not name: - return None, "`name` is required." if required else "`name` cannot be empty." + return None, "`name` is required" if required else "`name` cannot be empty" if len(name.encode("utf-8")) > 255: - return None, f"Chat name length is {len(name.encode('utf-8'))} which is larger than 255." + return None, f"chat name length is {len(name.encode('utf-8'))} which is larger than 255" return name, None @@ -366,7 +366,7 @@ async def create(): # Validate tenant_id should not be provided if req.get("tenant_id"): - return get_data_error_result(message="`tenant_id` must not be provided.") + return get_data_error_result(message="`tenant_id` must not be provided") # Validate name name, err = _validate_name(req.get("name"), required=True) @@ -424,7 +424,7 @@ async def create(): tenant_id=current_user.id, status=StatusEnum.VALID.value, ): - return get_data_error_result(message="Duplicated chat name in creating chat.") + return get_data_error_result(message="duplicated chat name in creating chat") req["id"] = get_uuid() req["tenant_id"] = current_user.id @@ -537,7 +537,7 @@ async def update_chat(chat_id): current_chat = current_chat.to_dict() if req.get("tenant_id"): - return get_data_error_result(message="`tenant_id` must not be provided.") + return get_data_error_result(message="`tenant_id` must not be provided") if "name" in req: name, err = _validate_name(req.get("name"), required=True) @@ -588,7 +588,7 @@ async def update_chat(chat_id): status=StatusEnum.VALID.value, ) ): - return get_data_error_result(message="Duplicated chat name.") + return get_data_error_result(message="duplicated chat name") if not DialogService.update_by_id(chat_id, req): return get_data_error_result(message="Chat not found!") @@ -676,7 +676,7 @@ async def patch_chat(chat_id): status=StatusEnum.VALID.value, ) ): - return get_data_error_result(message="Duplicated chat name.") + return get_data_error_result(message="duplicated chat name") if not DialogService.update_by_id(chat_id, req): return get_data_error_result(message="Failed to update chat.") diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 06538ceccd..54eeeb5f41 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -281,7 +281,7 @@ async def update_dataset(tenant_id: str, dataset_id: str, req: dict): :return: (success, result) or (success, error_message) """ if not req: - return False, "No properties were modified" + return False, "no properties were modified" kb = KnowledgebaseService.get_or_none(id=dataset_id, tenant_id=tenant_id) if kb is None: diff --git a/docs/references/http_api_reference.md b/docs/references/http_api_reference.md index 5fc8c8b231..f18f860488 100644 --- a/docs/references/http_api_reference.md +++ b/docs/references/http_api_reference.md @@ -3097,7 +3097,7 @@ Failure: ```json { "code": 102, - "message": "Duplicated chat name." + "message": "duplicated chat name" } ``` @@ -3234,7 +3234,7 @@ Failure: ```json { "code": 102, - "message": "Duplicated chat name." + "message": "duplicated chat name" } ``` diff --git a/internal/admin/handler.go b/internal/admin/handler.go index 0f8f108d6e..028d600455 100644 --- a/internal/admin/handler.go +++ b/internal/admin/handler.go @@ -1159,11 +1159,6 @@ func (h *Handler) Reports(c *gin.Context) { // Handle the heartbeat errCode, message := h.service.HandleHeartbeat(&req) - if errCode != common.CodeLicenseValid { - common.ErrorWithCode(c, errCode, message) - return - } - common.ErrorWithCode(c, errCode, message) } diff --git a/internal/admin/service.go b/internal/admin/service.go index 3e2a2a6eca..de48d3ae18 100644 --- a/internal/admin/service.go +++ b/internal/admin/service.go @@ -857,7 +857,7 @@ func (s *Service) ListUserAPITokens(ctx context.Context, username string) ([]map tenantID := userTenants[0].TenantID // 3. Get API tokens by tenant ID - tokens, err := s.apiTokenDAO.GetByTenantID(ctx, tenantID) + tokens, err := s.apiTokenDAO.GetByTenantID(ctx, dao.DB, tenantID) if err != nil { return nil, fmt.Errorf("failed to get API tokens: %w", err) } @@ -908,7 +908,7 @@ func (s *Service) GenerateUserAPIToken(ctx context.Context, username string) (ma } // 4. Save API token - if err = s.apiTokenDAO.Create(ctx, apiToken); err != nil { + if err = s.apiTokenDAO.Create(ctx, dao.DB, apiToken); err != nil { return nil, fmt.Errorf("failed to generate API key: %w", err) } @@ -940,7 +940,7 @@ func (s *Service) DeleteUserAPIToken(ctx context.Context, username, key string) tenantID := userTenants[0].TenantID // 3. Delete API token - rowsAffected, err := s.apiTokenDAO.DeleteByTenantIDAndToken(ctx, tenantID, key) + rowsAffected, err := s.apiTokenDAO.DeleteByTenantIDAndToken(ctx, dao.DB, tenantID, key) if err != nil { return fmt.Errorf("failed to delete API key: %w", err) } diff --git a/internal/dao/api_token.go b/internal/dao/api_token.go index 837235e428..cbc2c76d50 100644 --- a/internal/dao/api_token.go +++ b/internal/dao/api_token.go @@ -21,6 +21,8 @@ import ( "errors" "ragflow/internal/entity" + + "gorm.io/gorm" ) // APITokenDAO API token data access object @@ -32,27 +34,27 @@ func NewAPITokenDAO() *APITokenDAO { } // Create creates a new API token -func (dao *APITokenDAO) Create(ctx context.Context, apiToken *entity.APIToken) error { - return DB.WithContext(ctx).Create(apiToken).Error +func (dao *APITokenDAO) Create(ctx context.Context, db *gorm.DB, apiToken *entity.APIToken) error { + return db.WithContext(ctx).Create(apiToken).Error } // GetByTenantID gets API tokens by tenant ID -func (dao *APITokenDAO) GetByTenantID(ctx context.Context, tenantID string) ([]*entity.APIToken, error) { +func (dao *APITokenDAO) GetByTenantID(ctx context.Context, db *gorm.DB, tenantID string) ([]*entity.APIToken, error) { var tokens []*entity.APIToken - err := DB.WithContext(ctx).Where("tenant_id = ?", tenantID).Find(&tokens).Error + err := db.WithContext(ctx).Where("tenant_id = ?", tenantID).Find(&tokens).Error return tokens, err } // DeleteByTenantID deletes all API tokens by tenant ID (hard delete) -func (dao *APITokenDAO) DeleteByTenantID(ctx context.Context, tenantID string) (int64, error) { - result := DB.WithContext(ctx).Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.APIToken{}) +func (dao *APITokenDAO) DeleteByTenantID(ctx context.Context, db *gorm.DB, tenantID string) (int64, error) { + result := db.WithContext(ctx).Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.APIToken{}) return result.RowsAffected, result.Error } -// GetByToken gets API token by access key -func (dao *APITokenDAO) GetUserByAPIToken(ctx context.Context, token string) (*entity.APIToken, error) { +// GetUserByAPIToken gets user by API token +func (dao *APITokenDAO) GetUserByAPIToken(ctx context.Context, db *gorm.DB, token string) (*entity.APIToken, error) { var apiToken entity.APIToken - err := DB.WithContext(ctx).Where("token = ?", token).First(&apiToken).Error + err := db.WithContext(ctx).Where("token = ?", token).First(&apiToken).Error if err != nil { return nil, err } @@ -61,24 +63,24 @@ func (dao *APITokenDAO) GetUserByAPIToken(ctx context.Context, token string) (*e // GetByBeta gets API tokens by beta key (SDK/bot authorization token). // Mirrors Python's APIToken.query(beta=token), which returns a list. -func (dao *APITokenDAO) GetByBeta(ctx context.Context, beta string) ([]*entity.APIToken, error) { +func (dao *APITokenDAO) GetByBeta(ctx context.Context, db *gorm.DB, beta string) ([]*entity.APIToken, error) { var tokens []*entity.APIToken - err := DB.WithContext(ctx).Where("beta = ?", beta).Find(&tokens).Error + err := db.WithContext(ctx).Where("beta = ?", beta).Find(&tokens).Error return tokens, err } // DeleteByDialogIDs deletes API tokens by dialog IDs (hard delete) -func (dao *APITokenDAO) DeleteByDialogIDs(ctx context.Context, dialogIDs []string) (int64, error) { +func (dao *APITokenDAO) DeleteByDialogIDs(ctx context.Context, db *gorm.DB, dialogIDs []string) (int64, error) { if len(dialogIDs) == 0 { return 0, nil } - result := DB.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.APIToken{}) + result := db.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.APIToken{}) return result.RowsAffected, result.Error } // DeleteByTenantIDAndToken deletes a specific API token by tenant ID and token value -func (dao *APITokenDAO) DeleteByTenantIDAndToken(ctx context.Context, tenantID, token string) (int64, error) { - result := DB.WithContext(ctx).Unscoped().Where("tenant_id = ? AND token = ?", tenantID, token).Delete(&entity.APIToken{}) +func (dao *APITokenDAO) DeleteByTenantIDAndToken(ctx context.Context, db *gorm.DB, tenantID, token string) (int64, error) { + result := db.WithContext(ctx).Unscoped().Where("tenant_id = ? AND token = ?", tenantID, token).Delete(&entity.APIToken{}) return result.RowsAffected, result.Error } @@ -104,13 +106,13 @@ type ConversationStatsRow struct { // Create inserts a new api_4_conversation row. The caller is responsible // for setting ID, DialogID, UserID and the BaseModel time fields; the // DAO does not assign defaults because session creation paths in the -// Python agent API generate a uuid + tenant timestamp and rely on the +// Python agent API generate an uuid + tenant timestamp and rely on the // round-trip shape being byte-identical. -func (dao *API4ConversationDAO) Create(ctx context.Context, conv *entity.API4Conversation) error { +func (dao *API4ConversationDAO) Create(ctx context.Context, db *gorm.DB, conv *entity.API4Conversation) error { if conv == nil { return errors.New("api4 conversation: nil row") } - return DB.WithContext(ctx).Create(conv).Error + return db.WithContext(ctx).Create(conv).Error } // Update writes back an existing api_4_conversation row. The bot @@ -118,21 +120,21 @@ func (dao *API4ConversationDAO) Create(ctx context.Context, conv *entity.API4Con // turn so multi-turn chatbot sessions carry prior history into the next // LLM call. Matches the Python conversation_service.update pattern at // api/db/services/conversation_service.py:236 (async_iframe_completion). -func (dao *API4ConversationDAO) Update(ctx context.Context, conv *entity.API4Conversation) error { +func (dao *API4ConversationDAO) Update(ctx context.Context, db *gorm.DB, conv *entity.API4Conversation) error { if conv == nil { return errors.New("api4 conversation: nil row") } if conv.ID == "" { return errors.New("api4 conversation: empty id") } - return DB.WithContext(ctx).Save(conv).Error + return db.WithContext(ctx).Save(conv).Error } // Stats returns daily conversation aggregates for a tenant. -func (dao *API4ConversationDAO) Stats(ctx context.Context, tenantID, fromDate, toDate string, source *string) ([]ConversationStatsRow, error) { +func (dao *API4ConversationDAO) Stats(ctx context.Context, db *gorm.DB, tenantID, fromDate, toDate string, source *string) ([]ConversationStatsRow, error) { var rows []ConversationStatsRow dateExpr := "DATE_FORMAT(a.create_date, '%Y-%m-%d 00:00:00')" - db := DB.Table("api_4_conversation AS a"). + query := db.WithContext(ctx).Table("api_4_conversation AS a"). Select(` DATE_FORMAT(a.create_date, '%Y-%m-%d 00:00:00') AS dt, COUNT(a.id) AS pv, @@ -146,20 +148,20 @@ func (dao *API4ConversationDAO) Stats(ctx context.Context, tenantID, fromDate, t Where("a.create_date >= ? AND a.create_date <= ?", fromDate, toDate) if source == nil { - db = db.Where("a.source IS NULL") + query = query.Where("a.source IS NULL") } else { - db = db.Where("a.source = ?", *source) + query = query.Where("a.source = ?", *source) } - err := db.Group(dateExpr). + err := query.Group(dateExpr). Order(dateExpr). Scan(&rows).Error return rows, err } -func (dao *API4ConversationDAO) GetBySessionID(ctx context.Context, sessionID, agentID string) (*entity.API4Conversation, error) { +func (dao *API4ConversationDAO) GetBySessionID(ctx context.Context, db *gorm.DB, sessionID, agentID string) (*entity.API4Conversation, error) { var result entity.API4Conversation - tx := DB.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, agentID).Find(&result) + tx := db.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, agentID).Find(&result) if tx.Error != nil { return nil, tx.Error } @@ -170,23 +172,23 @@ func (dao *API4ConversationDAO) GetBySessionID(ctx context.Context, sessionID, a } // ListIDsByAgentID lists conversation IDs for one agent. -func (dao *API4ConversationDAO) ListIDsByAgentID(ctx context.Context, agentID string) ([]string, error) { +func (dao *API4ConversationDAO) ListIDsByAgentID(ctx context.Context, db *gorm.DB, agentID string) ([]string, error) { var ids []string - err := DB.WithContext(ctx).Model(&entity.API4Conversation{}).Where("dialog_id = ?", agentID).Pluck("id", &ids).Error + err := db.WithContext(ctx).Model(&entity.API4Conversation{}).Where("dialog_id = ?", agentID).Pluck("id", &ids).Error return ids, err } // DeleteBySessionIDAndAgentID deletes API4Conversations by sessionID and agentID -func (dao *API4ConversationDAO) DeleteBySessionIDAndAgentID(ctx context.Context, sessionID, agentID string) (int64, error) { - result := DB.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, agentID).Delete(&entity.API4Conversation{}) +func (dao *API4ConversationDAO) DeleteBySessionIDAndAgentID(ctx context.Context, db *gorm.DB, sessionID, agentID string) (int64, error) { + result := db.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, agentID).Delete(&entity.API4Conversation{}) return result.RowsAffected, result.Error } // DeleteByDialogIDs deletes API4Conversations by dialog IDs (hard delete) -func (dao *API4ConversationDAO) DeleteByDialogIDs(ctx context.Context, dialogIDs []string) (int64, error) { +func (dao *API4ConversationDAO) DeleteByDialogIDs(ctx context.Context, db *gorm.DB, dialogIDs []string) (int64, error) { if len(dialogIDs) == 0 { return 0, nil } - result := DB.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.API4Conversation{}) + result := db.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.API4Conversation{}) return result.RowsAffected, result.Error } diff --git a/internal/dao/api_token_beta_test.go b/internal/dao/api_token_beta_test.go index eb6264ef26..ab83e3aceb 100644 --- a/internal/dao/api_token_beta_test.go +++ b/internal/dao/api_token_beta_test.go @@ -56,7 +56,7 @@ func TestAPITokenDAOGetByBeta(t *testing.T) { } ctx := t.Context() - got, err := NewAPITokenDAO().GetByBeta(ctx, beta) + got, err := NewAPITokenDAO().GetByBeta(ctx, DB, beta) if err != nil { t.Fatalf("GetByBeta failed: %v", err) } diff --git a/internal/dao/api_token_test.go b/internal/dao/api_token_test.go index 78bebf72c4..eb273ee62d 100644 --- a/internal/dao/api_token_test.go +++ b/internal/dao/api_token_test.go @@ -64,7 +64,7 @@ func TestAPI4ConversationDAOGetBySessionID(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-1", "agent-1") ctx := t.Context() - session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, "session-1", "agent-1") + session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, DB, "session-1", "agent-1") if err != nil { t.Fatalf("GetBySessionID failed: %v", err) } @@ -86,7 +86,7 @@ func TestAPI4ConversationDAOGetBySessionIDWrongAgent(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-1", "agent-1") ctx := t.Context() - session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, "session-1", "agent-2") + session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, DB, "session-1", "agent-2") if err != nil { t.Fatalf("GetBySessionID failed: %v", err) } @@ -102,7 +102,7 @@ func TestAPI4ConversationDAOGetBySessionIDNoRows(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-1", "agent-1") ctx := t.Context() - session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, "missing-session", "agent-1") + session, err := NewAPI4ConversationDAO().GetBySessionID(ctx, DB, "missing-session", "agent-1") if err != nil { t.Fatalf("GetBySessionID failed: %v", err) } @@ -120,7 +120,7 @@ func TestAPI4ConversationDAOListIDsByAgentID(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-other", "agent-2") ctx := t.Context() - ids, err := NewAPI4ConversationDAO().ListIDsByAgentID(ctx, "agent-1") + ids, err := NewAPI4ConversationDAO().ListIDsByAgentID(ctx, DB, "agent-1") if err != nil { t.Fatalf("ListIDsByAgentID failed: %v", err) } @@ -144,7 +144,7 @@ func TestAPI4ConversationDAOListIDsByAgentIDNoRows(t *testing.T) { createAPI4ConversationForDAOTest(t, "session-1", "agent-1") ctx := t.Context() - ids, err := NewAPI4ConversationDAO().ListIDsByAgentID(ctx, "agent-2") + ids, err := NewAPI4ConversationDAO().ListIDsByAgentID(ctx, DB, "agent-2") if err != nil { t.Fatalf("ListIDsByAgentID failed: %v", err) } diff --git a/internal/dao/canvas_template.go b/internal/dao/canvas_template.go index 1c0a4ef907..7a0b93e4bd 100644 --- a/internal/dao/canvas_template.go +++ b/internal/dao/canvas_template.go @@ -19,6 +19,8 @@ package dao import ( "context" "ragflow/internal/entity" + + "gorm.io/gorm" ) // CanvasTemplateDAO data-access object for the canvas_template table. @@ -32,9 +34,9 @@ func NewCanvasTemplateDAO() *CanvasTemplateDAO { // GetAll returns every row in canvas_template ordered by create_time desc, so // templates appear newest first in the UI. Mirrors the Python // CanvasTemplateService.get_all() behaviour. -func (dao *CanvasTemplateDAO) GetAll(ctx context.Context) ([]*entity.CanvasTemplate, error) { +func (dao *CanvasTemplateDAO) GetAll(ctx context.Context, db *gorm.DB) ([]*entity.CanvasTemplate, error) { var templates []*entity.CanvasTemplate - if err := DB.WithContext(ctx).Order("create_time desc").Find(&templates).Error; err != nil { + if err := db.WithContext(ctx).Order("create_time desc").Find(&templates).Error; err != nil { return nil, err } return templates, nil diff --git a/internal/dao/canvas_template_seed.go b/internal/dao/canvas_template_seed.go index 9a3b48c65a..9bac66c47f 100644 --- a/internal/dao/canvas_template_seed.go +++ b/internal/dao/canvas_template_seed.go @@ -36,8 +36,8 @@ import ( // SeedCanvasTemplates seeds the canvas_template table from the built-in // agent/templates/*.json and internal/ingestion/pipeline/template/*.json files. -func SeedCanvasTemplates(ctx context.Context) error { - if err := addColumnIfNotExists(ctx, DB, "canvas_template", "parser_ids", "LONGTEXT NULL"); err != nil { +func SeedCanvasTemplates(ctx context.Context, db *gorm.DB) error { + if err := addColumnIfNotExists(ctx, db, "canvas_template", "parser_ids", "LONGTEXT NULL"); err != nil { return fmt.Errorf("failed to ensure canvas_template.parser_ids column: %w", err) } @@ -63,9 +63,9 @@ func SeedCanvasTemplates(ctx context.Context) error { return nil } - err := DB.Transaction(func(tx *gorm.DB) error { + err := db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { for _, tmpl := range allTemplates { - if err := tx.Clauses(clause.OnConflict{ + if err := tx.WithContext(ctx).Clauses(clause.OnConflict{ Columns: []clause.Column{{Name: "id"}}, DoUpdates: clause.AssignmentColumns([]string{ "avatar", "title", "description", "canvas_type", "canvas_types", "canvas_category", "dsl", @@ -74,7 +74,7 @@ func SeedCanvasTemplates(ctx context.Context) error { return fmt.Errorf("failed to save agent template %s: %w", tmpl.ID, err) } } - if err := tx.Where("id NOT IN ?", allIDs).Delete(&entity.CanvasTemplate{}).Error; err != nil { + if err := tx.WithContext(ctx).Where("id NOT IN ?", allIDs).Delete(&entity.CanvasTemplate{}).Error; err != nil { return fmt.Errorf("failed to remove stale agent templates: %w", err) } return nil @@ -136,7 +136,7 @@ func findIngestionTemplatesDir() string { return "" } -func seedCanvasTemplates(db *gorm.DB, dir string, entries []os.DirEntry) (int, error) { +func seedCanvasTemplates(ctx context.Context, db *gorm.DB, dir string, entries []os.DirEntry) (int, error) { templates := make([]*entity.CanvasTemplate, 0, len(entries)) ids := make([]string, 0, len(entries)) for _, entry := range entries { @@ -160,9 +160,9 @@ func seedCanvasTemplates(db *gorm.DB, dir string, entries []os.DirEntry) (int, e if len(templates) == 0 { return 0, nil } - err := db.Transaction(func(tx *gorm.DB) error { + err := db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { for _, tmpl := range templates { - if err := tx.Clauses(clause.OnConflict{ + if err := tx.WithContext(ctx).Clauses(clause.OnConflict{ Columns: []clause.Column{{Name: "id"}}, DoUpdates: clause.AssignmentColumns([]string{ "avatar", "title", "description", "canvas_type", "canvas_types", "canvas_category", "dsl", diff --git a/internal/dao/canvas_template_seed_test.go b/internal/dao/canvas_template_seed_test.go index d4e90a53ae..419469458b 100644 --- a/internal/dao/canvas_template_seed_test.go +++ b/internal/dao/canvas_template_seed_test.go @@ -49,18 +49,20 @@ func TestSeedCanvasTemplatesIsIdempotentAndRemovesStaleRows(t *testing.T) { seed := func() { t.Helper() - entries, err := os.ReadDir(dir) + ctx := t.Context() + var entries []os.DirEntry + entries, err = os.ReadDir(dir) if err != nil { t.Fatalf("read templates: %v", err) } - if _, err = seedCanvasTemplates(db, dir, entries); err != nil { + if _, err = seedCanvasTemplates(ctx, db, dir, entries); err != nil { t.Fatalf("seed templates: %v", err) } } seed() writeTemplate("kept.json", "kept", "updated") - if err := os.Remove(filepath.Join(dir, "removed.json")); err != nil { + if err = os.Remove(filepath.Join(dir, "removed.json")); err != nil { t.Fatalf("remove stale template file: %v", err) } seed() diff --git a/internal/dao/chat.go b/internal/dao/chat.go index a0369088a7..8f7cf8ae70 100644 --- a/internal/dao/chat.go +++ b/internal/dao/chat.go @@ -36,10 +36,10 @@ func NewChatDAO() *ChatDAO { } // ListByTenantID list chats by tenant ID -func (dao *ChatDAO) ListByTenantID(ctx context.Context, tenantID string, status string) ([]*entity.Chat, error) { +func (dao *ChatDAO) ListByTenantID(ctx context.Context, db *gorm.DB, tenantID string, status string) ([]*entity.Chat, error) { var chats []*entity.Chat - query := DB.WithContext(ctx).Model(&entity.Chat{}). + query := db.WithContext(ctx).Model(&entity.Chat{}). Where("tenant_id = ?", tenantID) if status != "" { @@ -55,12 +55,12 @@ func (dao *ChatDAO) ListByTenantID(ctx context.Context, tenantID string, status } // ListByTenantIDs list chats by tenant IDs with pagination and filtering -func (dao *ChatDAO) ListByTenantIDs(ctx context.Context, tenantIDs []string, userID string, page, pageSize int, orderby string, desc bool, keywords string) ([]*entity.ChatListItem, int64, error) { +func (dao *ChatDAO) ListByTenantIDs(ctx context.Context, db *gorm.DB, tenantIDs []string, userID string, page, pageSize int, orderby string, desc bool, keywords string) ([]*entity.ChatListItem, int64, error) { var chats []*entity.ChatListItem var total int64 // Build query with join to user table for nickname and avatar - query := DB.WithContext(ctx).Model(&entity.Chat{}). + query := db.WithContext(ctx).Model(&entity.Chat{}). Select(` dialog.*, user.nickname, @@ -107,11 +107,11 @@ func (dao *ChatDAO) ListByTenantIDs(ctx context.Context, tenantIDs []string, use } // ListByOwnerIDs list chats by owner IDs with filtering (manual pagination) -func (dao *ChatDAO) ListByOwnerIDs(ctx context.Context, ownerIDs []string, userID string, orderby string, desc bool, keywords string) ([]*entity.ChatListItem, int64, error) { +func (dao *ChatDAO) ListByOwnerIDs(ctx context.Context, db *gorm.DB, ownerIDs []string, userID string, orderby string, desc bool, keywords string) ([]*entity.ChatListItem, int64, error) { var chats []*entity.ChatListItem // Build query with join to user table - query := DB.WithContext(ctx).Model(&entity.Chat{}). + query := db.WithContext(ctx).Model(&entity.Chat{}). Select(` dialog.*, user.nickname, @@ -146,9 +146,9 @@ func (dao *ChatDAO) ListByOwnerIDs(ctx context.Context, ownerIDs []string, userI } // GetByID gets chat by ID -func (dao *ChatDAO) GetByID(ctx context.Context, id string) (*entity.Chat, error) { +func (dao *ChatDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.Chat, error) { var chat entity.Chat - err := DB.WithContext(ctx).Where("id = ?", id).First(&chat).Error + err := db.WithContext(ctx).Where("id = ?", id).First(&chat).Error if err != nil { return nil, err } @@ -156,9 +156,9 @@ func (dao *ChatDAO) GetByID(ctx context.Context, id string) (*entity.Chat, error } // GetByIDAndStatus gets chat by ID and status -func (dao *ChatDAO) GetByIDAndStatus(ctx context.Context, id string, status string) (*entity.Chat, error) { +func (dao *ChatDAO) GetByIDAndStatus(ctx context.Context, db *gorm.DB, id string, status string) (*entity.Chat, error) { var chat entity.Chat - err := DB.WithContext(ctx).Where("id = ? AND status = ?", id, status).First(&chat).Error + err := db.WithContext(ctx).Where("id = ? AND status = ?", id, status).First(&chat).Error if err != nil { return nil, err } @@ -166,32 +166,32 @@ func (dao *ChatDAO) GetByIDAndStatus(ctx context.Context, id string, status stri } // GetExistingNames gets existing dialog names for a tenant -func (dao *ChatDAO) GetExistingNames(ctx context.Context, tenantID string, status string) ([]string, error) { +func (dao *ChatDAO) GetExistingNames(ctx context.Context, db *gorm.DB, tenantID string, status string) ([]string, error) { var names []string - err := DB.WithContext(ctx).Model(&entity.Chat{}). + err := db.WithContext(ctx).Model(&entity.Chat{}). Where("tenant_id = ? AND status = ?", tenantID, status). Pluck("name", &names).Error return names, err } // ExistsByNameTenantStatus checks whether a chat with the given name exists. -func (dao *ChatDAO) ExistsByNameTenantStatus(ctx context.Context, name, tenantID, status string) (bool, error) { +func (dao *ChatDAO) ExistsByNameTenantStatus(ctx context.Context, db *gorm.DB, name, tenantID, status string) (bool, error) { var count int64 - err := DB.WithContext(ctx).Model(&entity.Chat{}). + err := db.WithContext(ctx).Model(&entity.Chat{}). Where("name = ? AND tenant_id = ? AND status = ?", name, tenantID, status). Count(&count).Error return count > 0, err } // Create creates a new chat/dialog -func (dao *ChatDAO) Create(ctx context.Context, chat *entity.Chat) error { +func (dao *ChatDAO) Create(ctx context.Context, db *gorm.DB, chat *entity.Chat) error { // Select("*") forces GORM to persist explicit zero values (e.g. similarity_threshold=0, // vector_similarity_weight=0, top_n=0) instead of substituting the column defaults. - return DB.WithContext(ctx).Select("*").Create(chat).Error + return db.WithContext(ctx).Select("*").Create(chat).Error } // UpdateByID updates a chat by ID -func (dao *ChatDAO) UpdateByID(ctx context.Context, id string, updates map[string]interface{}) error { +func (dao *ChatDAO) UpdateByID(ctx context.Context, db *gorm.DB, id string, updates map[string]interface{}) error { if updates == nil { updates = make(map[string]interface{}) } @@ -200,7 +200,7 @@ func (dao *ChatDAO) UpdateByID(ctx context.Context, id string, updates map[strin updates["update_time"] = now.UnixMilli() updates["update_date"] = now.Truncate(time.Second) - result := DB.WithContext(ctx).Session(&gorm.Session{SkipHooks: true}).Model(&entity.Chat{}).Where("id = ?", id).Updates(updates) + result := db.WithContext(ctx).Session(&gorm.Session{SkipHooks: true}).Model(&entity.Chat{}).Where("id = ?", id).Updates(updates) if result.Error != nil { return result.Error } @@ -217,13 +217,13 @@ func (dao *ChatDAO) UpdateByID(ctx context.Context, id string, updates map[strin } // UpdateManyByID updates multiple chats by ID (batch update) -func (dao *ChatDAO) UpdateManyByID(ctx context.Context, updates []map[string]interface{}) error { +func (dao *ChatDAO) UpdateManyByID(ctx context.Context, db *gorm.DB, updates []map[string]interface{}) error { if len(updates) == 0 { return nil } // Use transaction for batch update - tx := DB.WithContext(ctx).Begin() + tx := db.WithContext(ctx).Begin() if tx.Error != nil { return tx.Error } @@ -253,15 +253,15 @@ func (dao *ChatDAO) UpdateManyByID(ctx context.Context, updates []map[string]int } // DeleteByTenantID deletes all chats by tenant ID (hard delete) -func (dao *ChatDAO) DeleteByTenantID(ctx context.Context, tenantID string) (int64, error) { - result := DB.WithContext(ctx).Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Chat{}) +func (dao *ChatDAO) DeleteByTenantID(ctx context.Context, db *gorm.DB, tenantID string) (int64, error) { + result := db.WithContext(ctx).Unscoped().Where("tenant_id = ?", tenantID).Delete(&entity.Chat{}) return result.RowsAffected, result.Error } // GetAllDialogIDsByTenantID gets all dialog IDs by tenant ID -func (dao *ChatDAO) GetAllDialogIDsByTenantID(ctx context.Context, tenantID string) ([]string, error) { +func (dao *ChatDAO) GetAllDialogIDsByTenantID(ctx context.Context, db *gorm.DB, tenantID string) ([]string, error) { var dialogIDs []string - err := DB.WithContext(ctx).Model(&entity.Chat{}). + err := db.WithContext(ctx).Model(&entity.Chat{}). Where("tenant_id = ?", tenantID). Pluck("id", &dialogIDs).Error return dialogIDs, err @@ -270,8 +270,8 @@ func (dao *ChatDAO) GetAllDialogIDsByTenantID(ctx context.Context, tenantID stri // QueryByTenantIDAndID checks if a chat exists with given tenant_id and id // Reference: Python DialogService.query(tenant_id=tenant.tenant_id, id=chat_id, status=StatusEnum.VALID.value) // Used for permission verification in get_chat API -func (dao *ChatDAO) QueryByTenantIDAndID(ctx context.Context, tenantID string, chatID string, status string) ([]*entity.Chat, error) { +func (dao *ChatDAO) QueryByTenantIDAndID(ctx context.Context, db *gorm.DB, tenantID string, chatID string, status string) ([]*entity.Chat, error) { var chats []*entity.Chat - err := DB.WithContext(ctx).Where("tenant_id = ? AND id = ? AND status = ?", tenantID, chatID, status).Find(&chats).Error + err := db.WithContext(ctx).Where("tenant_id = ? AND id = ? AND status = ?", tenantID, chatID, status).Find(&chats).Error return chats, err } diff --git a/internal/dao/chat_channel.go b/internal/dao/chat_channel.go index 2fe6da29ab..4c67e9e04d 100644 --- a/internal/dao/chat_channel.go +++ b/internal/dao/chat_channel.go @@ -3,6 +3,8 @@ package dao import ( "context" "ragflow/internal/entity" + + "gorm.io/gorm" ) type ChatChannelDAO struct{} @@ -11,22 +13,22 @@ func NewChatChannel() *ChatChannelDAO { return &ChatChannelDAO{} } -func (dao *ChatChannelDAO) Create(ctx context.Context, channel *entity.ChatChannel) error { - return DB.WithContext(ctx).Create(channel).Error +func (dao *ChatChannelDAO) Create(ctx context.Context, db *gorm.DB, channel *entity.ChatChannel) error { + return db.WithContext(ctx).Create(channel).Error } -func (dao *ChatChannelDAO) GetByIDOnly(ctx context.Context, id string) (*entity.ChatChannel, error) { +func (dao *ChatChannelDAO) GetByIDOnly(ctx context.Context, db *gorm.DB, id string) (*entity.ChatChannel, error) { var channel entity.ChatChannel - err := DB.WithContext(ctx).Where("id = ?", id).First(&channel).Error + err := db.WithContext(ctx).Where("id = ?", id).First(&channel).Error if err != nil { return nil, err } return &channel, err } -func (dao *ChatChannelDAO) GetByID(ctx context.Context, id string, tenantID string) (*entity.ChatChannel, error) { +func (dao *ChatChannelDAO) GetByID(ctx context.Context, db *gorm.DB, id string, tenantID string) (*entity.ChatChannel, error) { var channel entity.ChatChannel - err := DB.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).First(&channel).Error + err := db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).First(&channel).Error if err != nil { return nil, err } @@ -34,20 +36,20 @@ func (dao *ChatChannelDAO) GetByID(ctx context.Context, id string, tenantID stri } // UpdateByID Update a single record by ID -func (dao *ChatChannelDAO) UpdateByID(ctx context.Context, id string, tenantID string, updates map[string]any) error { - return DB.WithContext(ctx).Model(&entity.ChatChannel{}).Where("id = ? AND tenant_id = ?", id, tenantID).Updates(updates).Error +func (dao *ChatChannelDAO) UpdateByID(ctx context.Context, db *gorm.DB, id string, tenantID string, updates map[string]any) error { + return db.WithContext(ctx).Model(&entity.ChatChannel{}).Where("id = ? AND tenant_id = ?", id, tenantID).Updates(updates).Error } // DeleteByID Delete a single record by ID -func (dao *ChatChannelDAO) DeleteByID(ctx context.Context, id string, tenantID string) error { - return DB.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&entity.ChatChannel{}).Error +func (dao *ChatChannelDAO) DeleteByID(ctx context.Context, db *gorm.DB, id string, tenantID string) error { + return db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&entity.ChatChannel{}).Error } // ListByTenantID List a single record by TenantID -func (dao *ChatChannelDAO) ListByTenantID(ctx context.Context, tenantID string) ([]*entity.ChatChannelListResponse, error) { +func (dao *ChatChannelDAO) ListByTenantID(ctx context.Context, db *gorm.DB, tenantID string) ([]*entity.ChatChannelListResponse, error) { results := make([]*entity.ChatChannelListResponse, 0) - err := DB.WithContext(ctx).Table("chat_channel"). + err := db.WithContext(ctx).Table("chat_channel"). Select("chat_channel.id, chat_channel.name, chat_channel.channel, chat_channel.chat_id, chat_channel.status, dialog.name as dialog_name"). Joins("LEFT JOIN dialog ON dialog.id = chat_channel.chat_id"). Where("chat_channel.tenant_id = ?", tenantID). diff --git a/internal/dao/chat_channel_test.go b/internal/dao/chat_channel_test.go index 074afa2fe3..1869cb6b7d 100644 --- a/internal/dao/chat_channel_test.go +++ b/internal/dao/chat_channel_test.go @@ -38,7 +38,7 @@ func setupChatChannelTestDB(t *testing.T) *gorm.DB { } // Migrate chat_channel and dialog (entity.Chat) tables - if err := db.AutoMigrate(&entity.ChatChannel{}, &entity.Chat{}); err != nil { + if err = db.AutoMigrate(&entity.ChatChannel{}, &entity.Chat{}); err != nil { t.Fatalf("failed to migrate: %v", err) } @@ -70,13 +70,13 @@ func TestChatChannelDAO_CRUD(t *testing.T) { ctx := t.Context() - err := dao.Create(ctx, cc) + err := dao.Create(ctx, db, cc) if err != nil { t.Fatalf("failed to create chat channel: %v", err) } // 2. Test GetByID - res, err := dao.GetByID(ctx, "chan-1", "tenant-1") + res, err := dao.GetByID(ctx, db, "chan-1", "tenant-1") if err != nil { t.Fatalf("failed to get chat channel: %v", err) } @@ -92,7 +92,7 @@ func TestChatChannelDAO_CRUD(t *testing.T) { } // 2b. Test tenant isolation for GetByID - _, err = dao.GetByID(ctx, "chan-1", "tenant-2") + _, err = dao.GetByID(ctx, db, "chan-1", "tenant-2") if err == nil { t.Fatalf("expected error (not found) when getting with wrong tenant, got nil") } @@ -102,12 +102,12 @@ func TestChatChannelDAO_CRUD(t *testing.T) { "name": "Updated WeCom Bot", } // Try updating with wrong tenant - err = dao.UpdateByID(ctx, "chan-1", "tenant-2", updates) + err = dao.UpdateByID(ctx, db, "chan-1", "tenant-2", updates) if err != nil { t.Fatalf("failed to run UpdateByID with wrong tenant: %v", err) } // Verify it was NOT updated (should still be "Test WeCom Bot" since wrong tenant was used) - res, err = dao.GetByID(ctx, "chan-1", "tenant-1") + res, err = dao.GetByID(ctx, db, "chan-1", "tenant-1") if err != nil { t.Fatalf("failed to get chat channel: %v", err) } @@ -116,12 +116,12 @@ func TestChatChannelDAO_CRUD(t *testing.T) { } // Update with correct tenant - err = dao.UpdateByID(ctx, "chan-1", "tenant-1", updates) + err = dao.UpdateByID(ctx, db, "chan-1", "tenant-1", updates) if err != nil { t.Fatalf("failed to update chat channel: %v", err) } - res, err = dao.GetByID(ctx, "chan-1", "tenant-1") + res, err = dao.GetByID(ctx, db, "chan-1", "tenant-1") if err != nil { t.Fatalf("failed to get updated chat channel: %v", err) } @@ -130,23 +130,23 @@ func TestChatChannelDAO_CRUD(t *testing.T) { } // 3b. Test DeleteByID with wrong tenant (should not delete) - err = dao.DeleteByID(ctx, "chan-1", "tenant-2") + err = dao.DeleteByID(ctx, db, "chan-1", "tenant-2") if err != nil { t.Fatalf("failed to delete with wrong tenant: %v", err) } // Verify it still exists for tenant-1 - _, err = dao.GetByID(ctx, "chan-1", "tenant-1") + _, err = dao.GetByID(ctx, db, "chan-1", "tenant-1") if err != nil { t.Fatalf("expected chat channel to still exist for tenant-1, got error: %v", err) } // 4. Test DeleteByID - err = dao.DeleteByID(ctx, "chan-1", "tenant-1") + err = dao.DeleteByID(ctx, db, "chan-1", "tenant-1") if err != nil { t.Fatalf("failed to delete chat channel: %v", err) } - _, err = dao.GetByID(ctx, "chan-1", "tenant-1") + _, err = dao.GetByID(ctx, db, "chan-1", "tenant-1") if err == nil { t.Fatalf("expected record not found error, got nil") } @@ -212,7 +212,7 @@ func TestChatChannelDAO_ListByTenantID(t *testing.T) { // Perform query ctx := t.Context() - list, err := dao.ListByTenantID(ctx, "tenant-1") + list, err := dao.ListByTenantID(ctx, db, "tenant-1") if err != nil { t.Fatalf("ListByTenantID failed: %v", err) } diff --git a/internal/dao/chat_session.go b/internal/dao/chat_session.go index 53cee49621..29c4533d13 100644 --- a/internal/dao/chat_session.go +++ b/internal/dao/chat_session.go @@ -52,9 +52,9 @@ func NewChatSessionDAO() *ChatSessionDAO { } // GetByID gets chat session by ID -func (dao *ChatSessionDAO) GetByID(ctx context.Context, id string) (*entity.ChatSession, error) { +func (dao *ChatSessionDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.ChatSession, error) { var conv entity.ChatSession - err := DB.WithContext(ctx).Where("id = ?", id).First(&conv).Error + err := db.WithContext(ctx).Where("id = ?", id).First(&conv).Error if err != nil { return nil, err } @@ -62,9 +62,9 @@ func (dao *ChatSessionDAO) GetByID(ctx context.Context, id string) (*entity.Chat } // GetBySessionIDAndChatID gets a chat session by session ID and chat ID. -func (dao *ChatSessionDAO) GetBySessionIDAndChatID(ctx context.Context, sessionID, chatID string) (*entity.ChatSession, error) { +func (dao *ChatSessionDAO) GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string) (*entity.ChatSession, error) { var conv entity.ChatSession - err := DB.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, chatID).First(&conv).Error + err := db.WithContext(ctx).Where("id = ? AND dialog_id = ?", sessionID, chatID).First(&conv).Error if err != nil { return nil, err } @@ -72,12 +72,12 @@ func (dao *ChatSessionDAO) GetBySessionIDAndChatID(ctx context.Context, sessionI } // Create creates a new chat session -func (dao *ChatSessionDAO) Create(ctx context.Context, conv *entity.ChatSession) error { - return DB.WithContext(ctx).Create(conv).Error +func (dao *ChatSessionDAO) Create(ctx context.Context, db *gorm.DB, conv *entity.ChatSession) error { + return db.WithContext(ctx).Create(conv).Error } // UpdateByID updates a chat session by ID -func (dao *ChatSessionDAO) UpdateByID(ctx context.Context, id string, updates map[string]interface{}) error { +func (dao *ChatSessionDAO) UpdateByID(ctx context.Context, db *gorm.DB, id string, updates map[string]interface{}) error { if updates == nil { updates = make(map[string]interface{}) } @@ -86,13 +86,13 @@ func (dao *ChatSessionDAO) UpdateByID(ctx context.Context, id string, updates ma updates["update_time"] = now.UnixMilli() updates["update_date"] = now.Truncate(time.Second) - result := DB.WithContext(ctx).Session(&gorm.Session{SkipHooks: true}).Model(&entity.ChatSession{}).Where("id = ?", id).Updates(updates) + result := db.WithContext(ctx).Session(&gorm.Session{SkipHooks: true}).Model(&entity.ChatSession{}).Where("id = ?", id).Updates(updates) if result.Error != nil { return result.Error } if result.RowsAffected == 0 { var count int64 - if err := DB.WithContext(ctx).Model(&entity.ChatSession{}).Where("id = ?", id).Count(&count).Error; err != nil { + if err := db.WithContext(ctx).Model(&entity.ChatSession{}).Where("id = ?", id).Count(&count).Error; err != nil { return err } if count == 0 { @@ -103,23 +103,23 @@ func (dao *ChatSessionDAO) UpdateByID(ctx context.Context, id string, updates ma } // DeleteByID deletes a chat session by ID (hard delete) -func (dao *ChatSessionDAO) DeleteByID(ctx context.Context, id string) error { - return DB.WithContext(ctx).Where("id = ?", id).Delete(&entity.ChatSession{}).Error +func (dao *ChatSessionDAO) DeleteByID(ctx context.Context, db *gorm.DB, id string) error { + return db.WithContext(ctx).Where("id = ?", id).Delete(&entity.ChatSession{}).Error } // ListByChatID lists chat sessions by chat ID -func (dao *ChatSessionDAO) ListByChatID(ctx context.Context, chatID string) ([]*entity.ChatSession, error) { +func (dao *ChatSessionDAO) ListByChatID(ctx context.Context, db *gorm.DB, chatID string) ([]*entity.ChatSession, error) { var convs []*entity.ChatSession - err := DB.WithContext(ctx).Where("dialog_id = ?", chatID). + err := db.WithContext(ctx).Where("dialog_id = ?", chatID). Order("create_time DESC"). Find(&convs).Error return convs, err } // CheckDialogExists checks if a dialog exists with given tenant_id and dialog_id -func (dao *ChatSessionDAO) CheckDialogExists(ctx context.Context, tenantID, chatID string) (bool, error) { +func (dao *ChatSessionDAO) CheckDialogExists(ctx context.Context, db *gorm.DB, tenantID, chatID string) (bool, error) { var count int64 - err := DB.WithContext(ctx).Model(&entity.Chat{}). + err := db.WithContext(ctx).Model(&entity.Chat{}). Where("tenant_id = ? AND id = ? AND status = ?", tenantID, chatID, common.StatusDialogValid). Count(&count).Error if err != nil { @@ -129,9 +129,9 @@ func (dao *ChatSessionDAO) CheckDialogExists(ctx context.Context, tenantID, chat } // GetDialogByID gets dialog by ID -func (dao *ChatSessionDAO) GetDialogByID(ctx context.Context, chatID string) (*entity.Chat, error) { +func (dao *ChatSessionDAO) GetDialogByID(ctx context.Context, db *gorm.DB, chatID string) (*entity.Chat, error) { var dialog entity.Chat - err := DB.WithContext(ctx).Where("id = ? AND status = ?", chatID, common.StatusDialogValid).First(&dialog).Error + err := db.WithContext(ctx).Where("id = ? AND status = ?", chatID, common.StatusDialogValid).First(&dialog).Error if err != nil { return nil, err } @@ -139,17 +139,17 @@ func (dao *ChatSessionDAO) GetDialogByID(ctx context.Context, chatID string) (*e } // DeleteByDialogIDs deletes chat sessions by dialog IDs (hard delete) -func (dao *ChatSessionDAO) DeleteByDialogIDs(ctx context.Context, dialogIDs []string) (int64, error) { +func (dao *ChatSessionDAO) DeleteByDialogIDs(ctx context.Context, db *gorm.DB, dialogIDs []string) (int64, error) { if len(dialogIDs) == 0 { return 0, nil } - result := DB.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.ChatSession{}) + result := db.WithContext(ctx).Unscoped().Where("dialog_id IN ?", dialogIDs).Delete(&entity.ChatSession{}) return result.RowsAffected, result.Error } -func (dao *ChatSessionDAO) ListAgentSessionNames(ctx context.Context, agentID, expUserID string) ([]map[string]interface{}, error) { +func (dao *ChatSessionDAO) ListAgentSessionNames(ctx context.Context, db *gorm.DB, agentID, expUserID string) ([]map[string]interface{}, error) { var rows []map[string]interface{} - err := DB.WithContext(ctx).Model(&entity.API4Conversation{}). + err := db.WithContext(ctx).Model(&entity.API4Conversation{}). Select("id", "name"). Where("dialog_id = ? AND exp_user_id = ?", agentID, expUserID). Order("create_date DESC"). @@ -184,8 +184,8 @@ func normalizeAgentSessionOrderBy(orderBy string) string { } } -func (dao *ChatSessionDAO) ListAgentSessions(ctx context.Context, params ListAgentSessionsParams) (int64, []*entity.API4Conversation, error) { - query := DB.WithContext(ctx).Model(&entity.API4Conversation{}).Where("dialog_id = ?", params.AgentID) +func (dao *ChatSessionDAO) ListAgentSessions(ctx context.Context, db *gorm.DB, params ListAgentSessionsParams) (int64, []*entity.API4Conversation, error) { + query := db.WithContext(ctx).Model(&entity.API4Conversation{}).Where("dialog_id = ?", params.AgentID) if !params.IncludeDSL { query = query.Omit("dsl") } diff --git a/internal/dao/chat_session_test.go b/internal/dao/chat_session_test.go index 114626e424..b57eabc05d 100644 --- a/internal/dao/chat_session_test.go +++ b/internal/dao/chat_session_test.go @@ -91,17 +91,16 @@ func createChatSessionForDAOTest(t *testing.T, db *gorm.DB, id, chatID, name str func TestChatSessionDAOUpdateByIDRefreshesTimestampsOnEmptyUpdate(t *testing.T) { db := setupChatSessionDAOTestDB(t) - pushDB(t, db) oldUpdateTime := int64(1000) createChatSessionForDAOTest(t, db, "session-1", "chat-1", "same", oldUpdateTime) ctx := t.Context() - if err := NewChatSessionDAO().UpdateByID(ctx, "session-1", map[string]interface{}{}); err != nil { + if err := NewChatSessionDAO().UpdateByID(ctx, db, "session-1", map[string]interface{}{}); err != nil { t.Fatalf("UpdateByID failed: %v", err) } - session, err := NewChatSessionDAO().GetByID(ctx, "session-1") + session, err := NewChatSessionDAO().GetByID(ctx, db, "session-1") if err != nil { t.Fatalf("GetByID failed: %v", err) } @@ -115,22 +114,20 @@ func TestChatSessionDAOUpdateByIDRefreshesTimestampsOnEmptyUpdate(t *testing.T) func TestChatSessionDAOUpdateByIDSameValueSucceeds(t *testing.T) { db := setupChatSessionDAOTestDB(t) - pushDB(t, db) createChatSessionForDAOTest(t, db, "session-1", "chat-1", "same", 1000) ctx := t.Context() - if err := NewChatSessionDAO().UpdateByID(ctx, "session-1", map[string]interface{}{"name": "same"}); err != nil { + if err := NewChatSessionDAO().UpdateByID(ctx, db, "session-1", map[string]interface{}{"name": "same"}); err != nil { t.Fatalf("UpdateByID failed: %v", err) } } func TestChatSessionDAOUpdateByIDMissingSession(t *testing.T) { db := setupChatSessionDAOTestDB(t) - pushDB(t, db) ctx := t.Context() - err := NewChatSessionDAO().UpdateByID(ctx, "missing", nil) + err := NewChatSessionDAO().UpdateByID(ctx, db, "missing", nil) if !errors.Is(err, gorm.ErrRecordNotFound) { t.Fatalf("expected ErrRecordNotFound, got %v", err) } @@ -138,7 +135,6 @@ func TestChatSessionDAOUpdateByIDMissingSession(t *testing.T) { func TestChatSessionDAOListAgentSessionsOrdersByUpdateTimeDesc(t *testing.T) { db := setupChatSessionDAOTestDB(t) - pushDB(t, db) createAgentSessionForDAOTest(t, db, "session-old", "agent-1", "user-1", 1000) createAgentSessionForDAOTest(t, db, "session-new", "agent-1", "user-1", 3000) @@ -146,7 +142,7 @@ func TestChatSessionDAOListAgentSessionsOrdersByUpdateTimeDesc(t *testing.T) { createAgentSessionForDAOTest(t, db, "session-other-agent", "agent-2", "user-1", 9999) ctx := t.Context() - total, sessions, err := NewChatSessionDAO().ListAgentSessions(ctx, ListAgentSessionsParams{ + total, sessions, err := NewChatSessionDAO().ListAgentSessions(ctx, db, ListAgentSessionsParams{ AgentID: "agent-1", Page: 1, PageSize: 10, @@ -177,7 +173,6 @@ func TestChatSessionDAOListAgentSessionsOrdersByUpdateTimeDesc(t *testing.T) { func TestChatSessionDAOListAgentSessionsFiltersAndPaginates(t *testing.T) { db := setupChatSessionDAOTestDB(t) - pushDB(t, db) createAgentSessionForDAOTest(t, db, "session-1", "agent-1", "user-1", 1000) createAgentSessionForDAOTest(t, db, "session-2", "agent-1", "user-1", 2000) @@ -185,7 +180,7 @@ func TestChatSessionDAOListAgentSessionsFiltersAndPaginates(t *testing.T) { createAgentSessionForDAOTest(t, db, "session-other-user", "agent-1", "user-2", 4000) ctx := t.Context() - total, sessions, err := NewChatSessionDAO().ListAgentSessions(ctx, ListAgentSessionsParams{ + total, sessions, err := NewChatSessionDAO().ListAgentSessions(ctx, db, ListAgentSessionsParams{ AgentID: "agent-1", UserID: "user-1", Page: 2, diff --git a/internal/dao/connector.go b/internal/dao/connector.go index 79e326d881..18b33d8ed3 100644 --- a/internal/dao/connector.go +++ b/internal/dao/connector.go @@ -17,6 +17,7 @@ package dao import ( + "context" "errors" "fmt" "ragflow/internal/entity" @@ -60,10 +61,10 @@ type ConnectorDatasetListItem struct { // ListByTenantID list connectors by tenant ID // Only selects id, name, source, status fields (matching Python implementation) -func (dao *ConnectorDAO) ListByTenantID(tenantID string) ([]*ConnectorListItem, error) { +func (dao *ConnectorDAO) ListByTenantID(ctx context.Context, db *gorm.DB, tenantID string) ([]*ConnectorListItem, error) { var connectors []*ConnectorListItem - err := DB.Model(&entity.Connector{}). + err := db.WithContext(ctx).Model(&entity.Connector{}). Select("id", "name", "source", "status"). Where("tenant_id = ?", tenantID). Find(&connectors).Error @@ -76,15 +77,15 @@ func (dao *ConnectorDAO) ListByTenantID(tenantID string) ([]*ConnectorListItem, } // ListByDatasetID lists connectors linked to a dataset. -func (dao *ConnectorDAO) ListByDatasetID(datasetID string) ([]*ConnectorDatasetListItem, error) { - return dao.ListByDatasetIDTx(DB, datasetID) +func (dao *ConnectorDAO) ListByDatasetID(ctx context.Context, db *gorm.DB, datasetID string) ([]*ConnectorDatasetListItem, error) { + return dao.ListByDatasetIDTx(ctx, db, datasetID) } // ListByDatasetIDTx lists connectors linked to a dataset using the caller's DB handle. -func (dao *ConnectorDAO) ListByDatasetIDTx(db *gorm.DB, datasetID string) ([]*ConnectorDatasetListItem, error) { +func (dao *ConnectorDAO) ListByDatasetIDTx(ctx context.Context, db *gorm.DB, datasetID string) ([]*ConnectorDatasetListItem, error) { var connectors []*ConnectorDatasetListItem - err := db.Model(&entity.Connector2Kb{}). + err := db.WithContext(ctx).Model(&entity.Connector2Kb{}). Select("connector.id, connector.source, connector.name, connector2kb.auto_parse, connector.status"). Joins("JOIN connector ON connector2kb.connector_id = connector.id"). Where("connector2kb.kb_id = ?", datasetID). @@ -104,20 +105,20 @@ type DatasetConnectorLink struct { } // LinkDatasetConnectors syncs connector2kb rows for a dataset. -func (dao *ConnectorDAO) LinkDatasetConnectors(kbID string, connectors []DatasetConnectorLink) error { - return DB.Transaction(func(tx *gorm.DB) error { +func (dao *ConnectorDAO) LinkDatasetConnectors(ctx context.Context, db *gorm.DB, kbID string, connectors []DatasetConnectorLink) error { + return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var kb entity.Knowledgebase - if err := tx.Select("tenant_id").Where("id = ? AND status = ?", kbID, string(entity.StatusValid)).First(&kb).Error; err != nil { + if err := db.WithContext(ctx).Select("tenant_id").Where("id = ? AND status = ?", kbID, string(entity.StatusValid)).First(&kb).Error; err != nil { return err } - return dao.LinkDatasetConnectorsTx(tx, kbID, kb.TenantID, connectors) + return dao.LinkDatasetConnectorsTx(ctx, tx, kbID, kb.TenantID, connectors) }) } // LinkDatasetConnectorsTx syncs connector2kb rows using the caller's transaction. -func (dao *ConnectorDAO) LinkDatasetConnectorsTx(tx *gorm.DB, kbID, tenantID string, connectors []DatasetConnectorLink) error { +func (dao *ConnectorDAO) LinkDatasetConnectorsTx(ctx context.Context, tx *gorm.DB, kbID, tenantID string, connectors []DatasetConnectorLink) error { var existing []entity.Connector2Kb - if err := tx.Where("kb_id = ?", kbID).Find(&existing).Error; err != nil { + if err := tx.WithContext(ctx).Where("kb_id = ?", kbID).Find(&existing).Error; err != nil { return err } @@ -135,7 +136,7 @@ func (dao *ConnectorDAO) LinkDatasetConnectorsTx(tx *gorm.DB, kbID, tenantID str } var fullConnector entity.Connector - if err := tx.Where("id = ? AND tenant_id = ?", connector.ID, tenantID).First(&fullConnector).Error; err != nil { + if err := tx.WithContext(ctx).Where("id = ? AND tenant_id = ?", connector.ID, tenantID).First(&fullConnector).Error; err != nil { if IsNotFoundErr(err) { return fmt.Errorf("%w: %s", errConnectorNotAccessible, connector.ID) } @@ -143,7 +144,7 @@ func (dao *ConnectorDAO) LinkDatasetConnectorsTx(tx *gorm.DB, kbID, tenantID str } if _, ok := oldConnectorIDs[connector.ID]; ok { - if err := tx.Model(&entity.Connector2Kb{}). + if err := tx.WithContext(ctx).Model(&entity.Connector2Kb{}). Where("connector_id = ? AND kb_id = ?", connector.ID, kbID). Update("auto_parse", autoParse).Error; err != nil { return err @@ -151,7 +152,7 @@ func (dao *ConnectorDAO) LinkDatasetConnectorsTx(tx *gorm.DB, kbID, tenantID str continue } - if err := tx.Create(&entity.Connector2Kb{ + if err := tx.WithContext(ctx).Create(&entity.Connector2Kb{ ID: utility.GenerateUUID(), ConnectorID: connector.ID, KbID: kbID, @@ -160,12 +161,12 @@ func (dao *ConnectorDAO) LinkDatasetConnectorsTx(tx *gorm.DB, kbID, tenantID str return err } - if err := scheduleConnectorTask(tx, connector.ID, kbID, connectorTaskTypeSync, true); err != nil { + if err := scheduleConnectorTask(ctx, tx, connector.ID, kbID, connectorTaskTypeSync, true); err != nil { return err } if connectorConfigBool(fullConnector.Config, "sync_deleted_files") { - if err := scheduleConnectorTask(tx, connector.ID, kbID, connectorTaskTypePrune, false); err != nil { + if err := scheduleConnectorTask(ctx, tx, connector.ID, kbID, connectorTaskTypePrune, false); err != nil { return err } } @@ -175,11 +176,11 @@ func (dao *ConnectorDAO) LinkDatasetConnectorsTx(tx *gorm.DB, kbID, tenantID str if _, ok := nextConnectorIDs[connectorID]; ok { continue } - if err := tx.Where("kb_id = ? AND connector_id = ?", kbID, connectorID). + if err := tx.WithContext(ctx).Where("kb_id = ? AND connector_id = ?", kbID, connectorID). Delete(&entity.Connector2Kb{}).Error; err != nil { return err } - if err := tx.Model(&entity.SyncLogs{}). + if err := tx.WithContext(ctx).Model(&entity.SyncLogs{}). Where("connector_id = ? AND kb_id = ? AND status IN ?", connectorID, kbID, []string{string(entity.TaskStatusSchedule), string(entity.TaskStatusRunning)}). Update("status", string(entity.TaskStatusCancel)).Error; err != nil { return err @@ -190,78 +191,78 @@ func (dao *ConnectorDAO) LinkDatasetConnectorsTx(tx *gorm.DB, kbID, tenantID str } // GetByID get connector by ID -func (dao *ConnectorDAO) GetByID(id string) (*entity.Connector, error) { +func (dao *ConnectorDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.Connector, error) { var connector entity.Connector - err := DB.Where("id = ?", id).First(&connector).Error + err := db.WithContext(ctx).Where("id = ?", id).First(&connector).Error if err != nil { return nil, err } return &connector, nil } -// Create create a new connector -func (dao *ConnectorDAO) Create(connector *entity.Connector) error { - return DB.Create(connector).Error +// Create a new connector +func (dao *ConnectorDAO) Create(ctx context.Context, db *gorm.DB, connector *entity.Connector) error { + return db.WithContext(ctx).Create(connector).Error } // UpdateByID update connector by ID -func (dao *ConnectorDAO) UpdateByID(id string, updates map[string]interface{}) error { - return DB.Model(&entity.Connector{}).Where("id = ?", id).Updates(updates).Error +func (dao *ConnectorDAO) UpdateByID(ctx context.Context, db *gorm.DB, id string, updates map[string]interface{}) error { + return db.WithContext(ctx).Model(&entity.Connector{}).Where("id = ?", id).Updates(updates).Error } // DeleteByID delete connector by ID -func (dao *ConnectorDAO) DeleteByID(id string) error { - return DB.Where("id = ?", id).Delete(&entity.Connector{}).Error +func (dao *ConnectorDAO) DeleteByID(ctx context.Context, db *gorm.DB, id string) error { + return db.WithContext(ctx).Where("id = ?", id).Delete(&entity.Connector{}).Error } // CancelRunningOrScheduledLogs marks active sync logs as canceled for a connector. -func (dao *ConnectorDAO) CancelRunningOrScheduledLogs(connectorID string) error { - return DB.Model(&entity.SyncLogs{}). +func (dao *ConnectorDAO) CancelRunningOrScheduledLogs(ctx context.Context, db *gorm.DB, connectorID string) error { + return db.WithContext(ctx).Model(&entity.SyncLogs{}). Where("connector_id = ? AND status IN ?", connectorID, []string{string(entity.TaskStatusSchedule), string(entity.TaskStatusRunning)}). Update("status", string(entity.TaskStatusCancel)).Error } // ScheduleConnectorTasks schedules sync and optional prune tasks for a connector. -func (dao *ConnectorDAO) ScheduleConnectorTasks(connectorID string) error { - return DB.Transaction(func(tx *gorm.DB) error { +func (dao *ConnectorDAO) ScheduleConnectorTasks(ctx context.Context, db *gorm.DB, connectorID string) error { + return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var connector entity.Connector - if err := tx.Where("id = ?", connectorID).First(&connector).Error; err != nil { + if err := tx.WithContext(ctx).Where("id = ?", connectorID).First(&connector).Error; err != nil { return err } var mappings []entity.Connector2Kb - if err := tx.Where("connector_id = ?", connectorID).Find(&mappings).Error; err != nil { + if err := tx.WithContext(ctx).Where("connector_id = ?", connectorID).Find(&mappings).Error; err != nil { return err } for _, mapping := range mappings { - if err := scheduleConnectorTask(tx, connectorID, mapping.KbID, connectorTaskTypeSync, false); err != nil { + if err := scheduleConnectorTask(ctx, tx, connectorID, mapping.KbID, connectorTaskTypeSync, false); err != nil { return err } if connectorConfigBool(connector.Config, "sync_deleted_files") { - if err := scheduleConnectorTask(tx, connectorID, mapping.KbID, connectorTaskTypePrune, false); err != nil { + if err := scheduleConnectorTask(ctx, tx, connectorID, mapping.KbID, connectorTaskTypePrune, false); err != nil { return err } } } - return tx.Model(&entity.Connector{}). + return tx.WithContext(ctx).Model(&entity.Connector{}). Where("id = ?", connectorID). Update("status", string(entity.TaskStatusSchedule)).Error }) } // ListDocumentsByKBAndSourceType lists connector documents in a dataset. -func (dao *ConnectorDAO) ListDocumentsByKBAndSourceType(kbID, sourceType string) ([]*entity.Document, error) { +func (dao *ConnectorDAO) ListDocumentsByKBAndSourceType(ctx context.Context, db *gorm.DB, kbID, sourceType string) ([]*entity.Document, error) { var documents []*entity.Document - err := DB.Where("kb_id = ? AND source_type = ?", kbID, sourceType).Find(&documents).Error + err := db.WithContext(ctx).Where("kb_id = ? AND source_type = ?", kbID, sourceType).Find(&documents).Error return documents, err } // RebuildConnector replaces old connector documents with scheduled sync tasks. -func (dao *ConnectorDAO) RebuildConnector(connector *entity.Connector, kbID string, documents []*entity.Document) error { - return DB.Transaction(func(tx *gorm.DB) error { - if err := tx.Where("connector_id = ? AND kb_id = ?", connector.ID, kbID).Delete(&entity.SyncLogs{}).Error; err != nil { +func (dao *ConnectorDAO) RebuildConnector(ctx context.Context, db *gorm.DB, connector *entity.Connector, kbID string, documents []*entity.Document) error { + return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if err := tx.WithContext(ctx).Where("connector_id = ? AND kb_id = ?", connector.ID, kbID).Delete(&entity.SyncLogs{}).Error; err != nil { return err } @@ -276,7 +277,7 @@ func (dao *ConnectorDAO) RebuildConnector(connector *entity.Connector, kbID stri } var mappings []entity.File2Document - if err := tx.Where("document_id IN ?", docIDs).Find(&mappings).Error; err != nil { + if err := tx.WithContext(ctx).Where("document_id IN ?", docIDs).Find(&mappings).Error; err != nil { return err } fileIDs := make([]string, 0, len(mappings)) @@ -292,23 +293,23 @@ func (dao *ConnectorDAO) RebuildConnector(connector *entity.Connector, kbID stri fileIDs = append(fileIDs, *mapping.FileID) } - if err := tx.Where("doc_id IN ?", docIDs).Delete(&entity.Task{}).Error; err != nil { + if err := tx.WithContext(ctx).Where("doc_id IN ?", docIDs).Delete(&entity.Task{}).Error; err != nil { return err } - if err := tx.Where("document_id IN ?", docIDs).Delete(&entity.File2Document{}).Error; err != nil { + if err := tx.WithContext(ctx).Where("document_id IN ?", docIDs).Delete(&entity.File2Document{}).Error; err != nil { return err } if len(fileIDs) > 0 { - if err := tx.Unscoped(). + if err := tx.WithContext(ctx).Unscoped(). Where("id IN ? AND source_type = ?", fileIDs, string(entity.FileSourceKnowledgebase)). Delete(&entity.File{}).Error; err != nil { return err } } - if err := tx.Where("id IN ?", docIDs).Delete(&entity.Document{}).Error; err != nil { + if err := tx.WithContext(ctx).Where("id IN ?", docIDs).Delete(&entity.Document{}).Error; err != nil { return err } - if err := tx.Model(&entity.Knowledgebase{}). + if err := tx.WithContext(ctx).Model(&entity.Knowledgebase{}). Where("id = ?", kbID). Updates(map[string]interface{}{ "doc_num": gorm.Expr("doc_num - ?", len(docIDs)), @@ -319,17 +320,17 @@ func (dao *ConnectorDAO) RebuildConnector(connector *entity.Connector, kbID stri } } - if err := tx.Model(&entity.Connector{}). + if err := tx.WithContext(ctx).Model(&entity.Connector{}). Where("id = ?", connector.ID). Update("status", string(entity.TaskStatusSchedule)).Error; err != nil { return err } - if err := createRebuildSyncLog(tx, connector.ID, kbID, connectorTaskTypeSync, true); err != nil { + if err := createRebuildSyncLog(ctx, tx, connector.ID, kbID, connectorTaskTypeSync, true); err != nil { return err } if syncDeletedFiles, _ := connector.Config["sync_deleted_files"].(bool); syncDeletedFiles { - if err := createRebuildSyncLog(tx, connector.ID, kbID, connectorTaskTypePrune, false); err != nil { + if err := createRebuildSyncLog(ctx, tx, connector.ID, kbID, connectorTaskTypePrune, false); err != nil { return err } } @@ -342,13 +343,13 @@ const ( connectorTaskTypePrune = "prune" ) -func createRebuildSyncLog(tx *gorm.DB, connectorID, kbID, taskType string, reindex bool) error { +func createRebuildSyncLog(ctx context.Context, tx *gorm.DB, connectorID, kbID, taskType string, reindex bool) error { fromBeginning := "0" if reindex { fromBeginning = "1" } now := time.Now().Local() - return tx.Create(&entity.SyncLogs{ + return tx.WithContext(ctx).Create(&entity.SyncLogs{ ID: utility.GenerateToken(), ConnectorID: connectorID, KbID: kbID, @@ -361,9 +362,9 @@ func createRebuildSyncLog(tx *gorm.DB, connectorID, kbID, taskType string, reind }).Error } -func scheduleConnectorTask(tx *gorm.DB, connectorID, kbID, taskType string, reindex bool) error { +func scheduleConnectorTask(ctx context.Context, tx *gorm.DB, connectorID, kbID, taskType string, reindex bool) error { var existing int64 - if err := tx.Model(&entity.SyncLogs{}). + if err := tx.WithContext(ctx).Model(&entity.SyncLogs{}). Where("connector_id = ? AND kb_id = ? AND task_type = ? AND status = ?", connectorID, kbID, taskType, string(entity.TaskStatusSchedule)). Count(&existing).Error; err != nil { return err @@ -376,7 +377,7 @@ func scheduleConnectorTask(tx *gorm.DB, connectorID, kbID, taskType string, rein var totalDocsIndexed int64 if taskType == connectorTaskTypeSync { var latest entity.SyncLogs - err := tx.Where("connector_id = ? AND kb_id = ? AND task_type = ? AND status = ?", connectorID, kbID, taskType, string(entity.TaskStatusDone)). + err := tx.WithContext(ctx).Where("connector_id = ? AND kb_id = ? AND task_type = ? AND status = ?", connectorID, kbID, taskType, string(entity.TaskStatusDone)). Order("update_time DESC"). First(&latest).Error if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { @@ -393,7 +394,7 @@ func scheduleConnectorTask(tx *gorm.DB, connectorID, kbID, taskType string, rein fromBeginning = "1" } now := time.Now().Local() - return tx.Create(&entity.SyncLogs{ + return tx.WithContext(ctx).Create(&entity.SyncLogs{ ID: utility.GenerateToken(), ConnectorID: connectorID, KbID: kbID, @@ -424,8 +425,8 @@ func connectorConfigBool(config map[string]interface{}, key string) bool { } // ListLogsByConnectorID lists sync logs for one connector with pagination. -func (dao *ConnectorDAO) ListLogsByConnectorID(connectorID string, offset, limit int) ([]*entity.ConnectorSyncLog, int64, error) { - baseQuery := DB.Model(&entity.SyncLogs{}). +func (dao *ConnectorDAO) ListLogsByConnectorID(ctx context.Context, db *gorm.DB, connectorID string, offset, limit int) ([]*entity.ConnectorSyncLog, int64, error) { + baseQuery := db.WithContext(ctx).Model(&entity.SyncLogs{}). Joins("JOIN connector ON sync_logs.connector_id = connector.id"). Joins("JOIN connector2kb ON sync_logs.connector_id = connector2kb.connector_id AND sync_logs.kb_id = connector2kb.kb_id"). Joins("JOIN knowledgebase ON sync_logs.kb_id = knowledgebase.id"). diff --git a/internal/dao/database.go b/internal/dao/database.go index b5dc6569cc..863af2df52 100644 --- a/internal/dao/database.go +++ b/internal/dao/database.go @@ -162,7 +162,7 @@ func InitDB(ctx context.Context, migrateDB bool) error { if migrateDB { common.Info("Migrating database schema...") for _, m := range dataModels { - if err = autoMigrateSafely(DB, m); err != nil { + if err = autoMigrateSafely(ctx, DB, m); err != nil { return fmt.Errorf("failed to migrate model %T: %w", m, err) } } @@ -176,7 +176,7 @@ func InitDB(ctx context.Context, migrateDB bool) error { // Seed built-in agent templates so the Go backend can serve the // "create agent from template" catalogue without relying on Python-side // initialization. - if err = SeedCanvasTemplates(ctx); err != nil { + if err = SeedCanvasTemplates(ctx, DB); err != nil { common.Warn("Failed to seed canvas templates", zap.Error(err)) } @@ -240,9 +240,9 @@ func findModelConfigDir() (string, error) { // autoMigrateSafely runs AutoMigrate and ignores duplicate index errors // This handles cases where indexes already exist (e.g., created by Python backend) -func autoMigrateSafely(db *gorm.DB, model interface{}) error { +func autoMigrateSafely(ctx context.Context, db *gorm.DB, model interface{}) error { //err := db.Debug().AutoMigrate(model) // to print debug info - err := db.AutoMigrate(model) + err := db.WithContext(ctx).AutoMigrate(model) if err == nil { return nil } diff --git a/internal/handler/connector.go b/internal/handler/connector.go index 99d5a532ab..2fbde2534d 100644 --- a/internal/handler/connector.go +++ b/internal/handler/connector.go @@ -17,6 +17,7 @@ package handler import ( + "context" "encoding/json" "errors" "net/http" @@ -31,20 +32,20 @@ import ( ) type connectorServiceIface interface { - ListConnectors(userID string) (*service.ListConnectorsResponse, error) - CreateConnector(userID string, req *service.CreateConnectorRequest) (*entity.Connector, error) - GetConnector(connectorID, userID string) (*entity.Connector, common.ErrorCode, error) - ListLog(connectorID, userID string, page, pageSize int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) - DeleteConnector(connectorID, userID string) (bool, common.ErrorCode, error) - RebuildConnector(connectorID, userID, kbID string) (bool, common.ErrorCode, error) - TestConnector(connectorID, userID string) error - UpdateConnector(connectorID, userID string, req *service.UpdateConnectorRequest) (*entity.Connector, common.ErrorCode, error) - StartGoogleWebOAuth(userID, source string, req *service.StartGoogleWebOAuthRequest) (*service.StartGoogleWebOAuthResponse, common.ErrorCode, error) - GoogleWebOAuthCallback(source, stateID, oauthError, errorDescription, code string) string - PollGoogleWebOAuthResult(userID, source string, req *service.PollGoogleWebOAuthResultRequest) (*service.PollGoogleWebOAuthResultResponse, common.ErrorCode, error) - StartBoxWebOAuth(userID string, req *service.StartBoxWebOAuthRequest) (*service.StartBoxWebOAuthResponse, common.ErrorCode, error) - BoxWebOAuthCallback(flowID string, oauthError string, errorDescription string, code string) string - PollBoxWebOAuthResult(userID string, req *service.PollBoxWebOAuthResultRequest) (*service.PollBoxWebOAuthResultResponse, common.ErrorCode, error) + ListConnectors(ctx context.Context, userID string) (*service.ListConnectorsResponse, error) + CreateConnector(ctx context.Context, userID string, req *service.CreateConnectorRequest) (*entity.Connector, error) + GetConnector(ctx context.Context, connectorID, userID string) (*entity.Connector, common.ErrorCode, error) + ListLog(ctx context.Context, connectorID, userID string, page, pageSize int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) + DeleteConnector(ctx context.Context, connectorID, userID string) (bool, common.ErrorCode, error) + RebuildConnector(ctx context.Context, connectorID, userID, kbID string) (bool, common.ErrorCode, error) + TestConnector(ctx context.Context, connectorID, userID string) error + UpdateConnector(ctx context.Context, connectorID, userID string, req *service.UpdateConnectorRequest) (*entity.Connector, common.ErrorCode, error) + StartGoogleWebOAuth(ctx context.Context, userID, source string, req *service.StartGoogleWebOAuthRequest) (*service.StartGoogleWebOAuthResponse, common.ErrorCode, error) + GoogleWebOAuthCallback(ctx context.Context, source, stateID, oauthError, errorDescription, code string) string + PollGoogleWebOAuthResult(ctx context.Context, userID, source string, req *service.PollGoogleWebOAuthResultRequest) (*service.PollGoogleWebOAuthResultResponse, common.ErrorCode, error) + StartBoxWebOAuth(ctx context.Context, userID string, req *service.StartBoxWebOAuthRequest) (*service.StartBoxWebOAuthResponse, common.ErrorCode, error) + BoxWebOAuthCallback(ctx context.Context, flowID string, oauthError string, errorDescription string, code string) string + PollBoxWebOAuthResult(ctx context.Context, userID string, req *service.PollBoxWebOAuthResultRequest) (*service.PollBoxWebOAuthResultResponse, common.ErrorCode, error) } // ConnectorHandler connector handler @@ -76,9 +77,10 @@ func (h *ConnectorHandler) ListConnectors(c *gin.Context) { return } userID := user.ID + ctx := c.Request.Context() // List connectors - result, err := h.connectorService.ListConnectors(userID) + result, err := h.connectorService.ListConnectors(ctx, userID) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, 500, nil, err.Error()) return @@ -120,8 +122,9 @@ func (h *ConnectorHandler) GetConnector(c *gin.Context) { common.ErrorWithCode(c, errorCode, errorMessage) return } + ctx := c.Request.Context() - connector, code, err := h.connectorService.GetConnector(c.Param("connector_id"), user.ID) + connector, code, err := h.connectorService.GetConnector(ctx, c.Param("connector_id"), user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -144,7 +147,9 @@ func (h *ConnectorHandler) UpdateConnector(c *gin.Context) { return } - connector, code, err := h.connectorService.UpdateConnector(c.Param("connector_id"), user.ID, req) + ctx := c.Request.Context() + + connector, code, err := h.connectorService.UpdateConnector(ctx, c.Param("connector_id"), user.ID, req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -215,7 +220,9 @@ func (h *ConnectorHandler) ListLogs(c *gin.Context) { pageSize = parsedPageSize } - logs, total, code, err := h.connectorService.ListLog(c.Param("connector_id"), user.ID, page, pageSize) + ctx := c.Request.Context() + + logs, total, code, err := h.connectorService.ListLog(ctx, c.Param("connector_id"), user.ID, page, pageSize) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -261,7 +268,9 @@ func (h *ConnectorHandler) CreateConnector(c *gin.Context) { return } - connector, err := h.connectorService.CreateConnector(user.ID, &req) + ctx := c.Request.Context() + + connector, err := h.connectorService.CreateConnector(ctx, user.ID, &req) if err != nil { common.ResponseWithHttpCodeData(c, http.StatusInternalServerError, common.CodeServerError, nil, err.Error()) return @@ -290,7 +299,9 @@ func (h *ConnectorHandler) TestConnector(c *gin.Context) { return } - err := h.connectorService.TestConnector(connectorID, user.ID) + ctx := c.Request.Context() + + err := h.connectorService.TestConnector(ctx, connectorID, user.ID) if errors.Is(err, service.ErrConnectorTestUnsupported) { connectorErrorResponse(c, err) return @@ -319,7 +330,9 @@ func (h *ConnectorHandler) DeleteConnector(c *gin.Context) { return } - ok, code, err := h.connectorService.DeleteConnector(c.Param("connector_id"), user.ID) + ctx := c.Request.Context() + + ok, code, err := h.connectorService.DeleteConnector(ctx, c.Param("connector_id"), user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -357,7 +370,8 @@ func (h *ConnectorHandler) RebuildConnector(c *gin.Context) { return } - ok, code, err := h.connectorService.RebuildConnector(c.Param("connector_id"), user.ID, req.KbID) + ctx := c.Request.Context() + ok, code, err := h.connectorService.RebuildConnector(ctx, c.Param("connector_id"), user.ID, req.KbID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -379,7 +393,9 @@ func (h *ConnectorHandler) StartGoogleWebOAuth(c *gin.Context) { return } - data, code, err := h.connectorService.StartGoogleWebOAuth(user.ID, c.DefaultQuery("type", "google-drive"), &req) + ctx := c.Request.Context() + + data, code, err := h.connectorService.StartGoogleWebOAuth(ctx, user.ID, c.DefaultQuery("type", "google-drive"), &req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -401,7 +417,9 @@ func (h *ConnectorHandler) PollGoogleWebOAuthResult(c *gin.Context) { return } - data, code, err := h.connectorService.PollGoogleWebOAuthResult(user.ID, c.Query("type"), &req) + ctx := c.Request.Context() + + data, code, err := h.connectorService.PollGoogleWebOAuthResult(ctx, user.ID, c.Query("type"), &req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -423,7 +441,9 @@ func (h *ConnectorHandler) GmailWebOAuthCallback(c *gin.Context) { } func (h *ConnectorHandler) googleWebOAuthCallback(c *gin.Context, source string) { - html := h.connectorService.GoogleWebOAuthCallback( + ctx := c.Request.Context() + + html := h.connectorService.GoogleWebOAuthCallback(ctx, source, c.Query("state"), c.Query("error"), @@ -444,7 +464,9 @@ func (h *ConnectorHandler) StartBoxWebOAuth(c *gin.Context) { common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) return } - resp, code, err := h.connectorService.StartBoxWebOAuth(user.ID, &req) + ctx := c.Request.Context() + + resp, code, err := h.connectorService.StartBoxWebOAuth(ctx, user.ID, &req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -458,7 +480,9 @@ func (h *ConnectorHandler) BoxWebOAuthCallback(c *gin.Context) { errorDescription := c.Query("error_description") code := c.Query("code") - html := h.connectorService.BoxWebOAuthCallback(flowID, oauthError, errorDescription, code) + ctx := c.Request.Context() + + html := h.connectorService.BoxWebOAuthCallback(ctx, flowID, oauthError, errorDescription, code) c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(html)) } @@ -474,7 +498,9 @@ func (h *ConnectorHandler) PollBoxWebOAuthResult(c *gin.Context) { common.ErrorWithCode(c, common.CodeBadRequest, err.Error()) return } - resp, code, err := h.connectorService.PollBoxWebOAuthResult(user.ID, &req) + ctx := c.Request.Context() + + resp, code, err := h.connectorService.PollBoxWebOAuthResult(ctx, user.ID, &req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return diff --git a/internal/handler/connector_test.go b/internal/handler/connector_test.go index 602f2faaaf..506e6c85b2 100644 --- a/internal/handler/connector_test.go +++ b/internal/handler/connector_test.go @@ -1,6 +1,7 @@ package handler import ( + "context" "encoding/json" "fmt" "net/http" @@ -25,54 +26,54 @@ type fakeConnectorService struct { html string } -func (s fakeConnectorService) ListConnectors(string) (*service.ListConnectorsResponse, error) { +func (s fakeConnectorService) ListConnectors(context.Context, string) (*service.ListConnectorsResponse, error) { return &service.ListConnectorsResponse{}, nil } -func (s fakeConnectorService) TestConnector(string, string) error { +func (s fakeConnectorService) TestConnector(context.Context, string, string) error { return s.err } -func (s fakeConnectorService) CreateConnector(string, *service.CreateConnectorRequest) (*entity.Connector, error) { +func (s fakeConnectorService) CreateConnector(context.Context, string, *service.CreateConnectorRequest) (*entity.Connector, error) { if s.err != nil { return nil, s.err } return s.connector, nil } -func (s fakeConnectorService) GetConnector(string, string) (*entity.Connector, common.ErrorCode, error) { +func (s fakeConnectorService) GetConnector(context.Context, string, string) (*entity.Connector, common.ErrorCode, error) { if s.err != nil { return nil, s.code, s.err } return s.connector, common.CodeSuccess, nil } -func (s fakeConnectorService) UpdateConnector(string, string, *service.UpdateConnectorRequest) (*entity.Connector, common.ErrorCode, error) { +func (s fakeConnectorService) UpdateConnector(context.Context, string, string, *service.UpdateConnectorRequest) (*entity.Connector, common.ErrorCode, error) { if s.err != nil { return nil, s.code, s.err } return s.connector, common.CodeSuccess, nil } -func (s fakeConnectorService) StartGoogleWebOAuth(string, string, *service.StartGoogleWebOAuthRequest) (*service.StartGoogleWebOAuthResponse, common.ErrorCode, error) { +func (s fakeConnectorService) StartGoogleWebOAuth(context.Context, string, string, *service.StartGoogleWebOAuthRequest) (*service.StartGoogleWebOAuthResponse, common.ErrorCode, error) { if s.err != nil { return nil, s.code, s.err } return &service.StartGoogleWebOAuthResponse{}, common.CodeSuccess, nil } -func (s fakeConnectorService) GoogleWebOAuthCallback(string, string, string, string, string) string { +func (s fakeConnectorService) GoogleWebOAuthCallback(context.Context, string, string, string, string, string) string { return "" } -func (s fakeConnectorService) PollGoogleWebOAuthResult(string, string, *service.PollGoogleWebOAuthResultRequest) (*service.PollGoogleWebOAuthResultResponse, common.ErrorCode, error) { +func (s fakeConnectorService) PollGoogleWebOAuthResult(context.Context, string, string, *service.PollGoogleWebOAuthResultRequest) (*service.PollGoogleWebOAuthResultResponse, common.ErrorCode, error) { if s.err != nil { return nil, s.code, s.err } return &service.PollGoogleWebOAuthResultResponse{}, common.CodeSuccess, nil } -func (s fakeConnectorService) StartBoxWebOAuth(string, *service.StartBoxWebOAuthRequest) (*service.StartBoxWebOAuthResponse, common.ErrorCode, error) { +func (s fakeConnectorService) StartBoxWebOAuth(context.Context, string, *service.StartBoxWebOAuthRequest) (*service.StartBoxWebOAuthResponse, common.ErrorCode, error) { if s.err != nil { return nil, s.code, s.err } @@ -83,35 +84,35 @@ func (s fakeConnectorService) StartBoxWebOAuth(string, *service.StartBoxWebOAuth }, common.CodeSuccess, nil } -func (s fakeConnectorService) BoxWebOAuthCallback(string, string, string, string) string { +func (s fakeConnectorService) BoxWebOAuthCallback(context.Context, string, string, string, string) string { if s.html != "" { return s.html } return "box" } -func (s fakeConnectorService) PollBoxWebOAuthResult(string, *service.PollBoxWebOAuthResultRequest) (*service.PollBoxWebOAuthResultResponse, common.ErrorCode, error) { +func (s fakeConnectorService) PollBoxWebOAuthResult(context.Context, string, *service.PollBoxWebOAuthResultRequest) (*service.PollBoxWebOAuthResultResponse, common.ErrorCode, error) { if s.err != nil { return nil, s.code, s.err } return &service.PollBoxWebOAuthResultResponse{}, common.CodeSuccess, nil } -func (s fakeConnectorService) ListLog(string, string, int, int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) { +func (s fakeConnectorService) ListLog(context.Context, string, string, int, int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) { if s.err != nil { return nil, 0, s.code, s.err } return s.logs, s.total, common.CodeSuccess, nil } -func (s fakeConnectorService) DeleteConnector(string, string) (bool, common.ErrorCode, error) { +func (s fakeConnectorService) DeleteConnector(context.Context, string, string) (bool, common.ErrorCode, error) { if s.err != nil { return false, s.code, s.err } return true, common.CodeSuccess, nil } -func (s fakeConnectorService) RebuildConnector(string, string, string) (bool, common.ErrorCode, error) { +func (s fakeConnectorService) RebuildConnector(context.Context, string, string, string) (bool, common.ErrorCode, error) { if s.err != nil { return false, s.code, s.err } diff --git a/internal/handler/dataset.go b/internal/handler/dataset.go index 1ebbd2b432..28dbb29670 100644 --- a/internal/handler/dataset.go +++ b/internal/handler/dataset.go @@ -55,7 +55,7 @@ type listDatasetsExt struct { ParserID string `json:"parser_id,omitempty"` } -// NewDatasetsHandler creates a new datasets handler. +// NewDatasetsHandler creates a new datasets' handler. func NewDatasetsHandler(datasetsService *dataset.DatasetService, metadataService *service.MetadataService) *DatasetsHandler { h := &DatasetsHandler{ datasetsService: datasetsService, @@ -171,8 +171,10 @@ func (h *DatasetsHandler) GetDataset(c *gin.Context) { return } + ctx := c.Request.Context() + datasetID := c.Param("dataset_id") - result, code, err := h.datasetsService.GetDataset(datasetID, user.ID) + result, code, err := h.datasetsService.GetDataset(ctx, datasetID, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -208,19 +210,21 @@ func (h *DatasetsHandler) UpdateDataset(c *gin.Context) { } var req service.UpdateDatasetRequest - if err := json.Unmarshal(bodyBytes, &req); err != nil { + if err = json.Unmarshal(bodyBytes, &req); err != nil { common.ResponseWithCodeData(c, common.CodeDataError, nil, err.Error()) return } // Detect an explicitly provided parser_config key (even {} or null) so it is not - // rejected as "No properties were modified", mirroring the Python contract. + // rejected as "no properties were modified", mirroring the Python contract. var providedFields map[string]json.RawMessage - if err := json.Unmarshal(bodyBytes, &providedFields); err == nil { + if err = json.Unmarshal(bodyBytes, &providedFields); err == nil { _, req.ParserConfigProvided = providedFields["parser_config"] } - result, code, err := h.datasetsService.UpdateDataset(datasetID, userID, req) + ctx := c.Request.Context() + + result, code, err := h.datasetsService.UpdateDataset(ctx, datasetID, userID, req) if err != nil { common.ErrorWithCode(c, code, err.Error()) return @@ -407,13 +411,14 @@ func (h *DatasetsHandler) GetKnowledgeGraph(c *gin.Context) { return } - dataset, code, err := h.datasetsService.GetDataset(datasetID, user.ID) + ctx := c.Request.Context() + datasetInstance, code, err := h.datasetsService.GetDataset(ctx, datasetID, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return } - tenantID, _ := dataset["tenant_id"].(string) + tenantID, _ := datasetInstance["tenant_id"].(string) if tenantID == "" { common.ResponseWithCodeData(c, common.CodeDataError, nil, "tenant_id is required") return @@ -470,7 +475,7 @@ func (h *DatasetsHandler) GetKnowledgeGraph(c *gin.Context) { } var graphData map[string]interface{} - if err := json.Unmarshal([]byte(contentWithWeight), &graphData); err != nil { + if err = json.Unmarshal([]byte(contentWithWeight), &graphData); err != nil { common.SuccessWithData(c, result, "success") return } @@ -577,13 +582,14 @@ func (h *DatasetsHandler) DeleteKnowledgeGraph(c *gin.Context) { return } - dataset, code, err := h.datasetsService.GetDataset(datasetID, user.ID) + ctx := c.Request.Context() + datasetInstance, code, err := h.datasetsService.GetDataset(ctx, datasetID, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return } - tenantID, _ := dataset["tenant_id"].(string) + tenantID, _ := datasetInstance["tenant_id"].(string) if tenantID == "" { common.ResponseWithCodeData(c, common.CodeDataError, nil, "tenant_id is required") return @@ -611,8 +617,6 @@ func (h *DatasetsHandler) DeleteKnowledgeGraph(c *gin.Context) { // @Summary Remove Tags // @Description Remove tags from a dataset // @Tags datasets -// @Accept json -// @Produce json // @Security ApiKeyAuth // @Param dataset_id path string true "Dataset ID" // @Param request body object{tags []string} true "tags to remove" @@ -631,13 +635,14 @@ func (h *DatasetsHandler) RemoveTags(c *gin.Context) { return } - dataset, code, err := h.datasetsService.GetDataset(datasetID, user.ID) + ctx := c.Request.Context() + datasetInstance, code, err := h.datasetsService.GetDataset(ctx, datasetID, user.ID) if err != nil { common.ErrorWithCode(c, code, err.Error()) return } - tenantID, _ := dataset["tenant_id"].(string) + tenantID, _ := datasetInstance["tenant_id"].(string) if tenantID == "" { common.ResponseWithCodeData(c, common.CodeDataError, nil, "tenant_id is required") return @@ -958,8 +963,6 @@ func (h *DatasetsHandler) UpdateDocumentMetadataConfig(c *gin.Context) { // @Summary Search Datasets // @Description Search for relevant chunks across one or more datasets based on a question // @Tags datasets -// @Accept json -// @Produce json // @Param request body service.SearchDatasetsRequest true "search parameters" // @Success 200 {object} map[string]interface{} // @Router /api/v1/datasets/search [post] @@ -1034,8 +1037,6 @@ func (h *DatasetsHandler) SearchDatasets(c *gin.Context) { // @Summary Search Dataset // @Description Search for relevant chunks within one dataset based on a question // @Tags datasets -// @Accept json -// @Produce json // @Param dataset_id path string true "dataset id" // @Param request body service.SearchDatasetRequest true "search parameters" // @Success 200 {object} map[string]interface{} diff --git a/internal/parser/parser/pdf_parser_nocgo.go b/internal/parser/parser/pdf_parser_nocgo.go index 48fb59165d..1a4f808abe 100644 --- a/internal/parser/parser/pdf_parser_nocgo.go +++ b/internal/parser/parser/pdf_parser_nocgo.go @@ -6,7 +6,7 @@ import ( "fmt" ) -func (p *PDFParser) ParseWithResult(filename string, data []byte) ParseResult { +func (p *PDFParser) ParseWithResult(ctx context.Context, filename string, data []byte) ParseResult { if err := p.validateParseMethod(); err != nil { return ParseResult{Err: err} } diff --git a/internal/service/agent.go b/internal/service/agent.go index b582383533..b1e8d59e6a 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -356,7 +356,7 @@ func NewAgentServiceWithOptions( // agent_api.list_agent_template, which iterates CanvasTemplateService.get_all() // and serialises each row. func (s *AgentService) ListTemplates(ctx context.Context) ([]*entity.CanvasTemplate, error) { - return s.canvasTemplateDAO.GetAll(ctx) + return s.canvasTemplateDAO.GetAll(ctx, dao.DB) } // AgentItem is one entry in the list response. @@ -668,7 +668,7 @@ func (s *AgentService) UpdateAgent(ctx context.Context, userID, canvasID string, if title, ok := updatedAgentTitle(canvasInstance, updates); ok { return agentTitleAlreadyExistsError(title) } - return errors.New("Agent title already exists.") + return errors.New("agent title already exists") } return fmt.Errorf("update agent %s: %w", canvasID, err) } @@ -1020,7 +1020,7 @@ func (s *AgentService) RunAgent(ctx context.Context, userID, canvasID, sessionID dsl = normalisedDSLForRun(versionRow) } if sessionID != "" && s.api4ConversationDAO != nil { - session, sessionErr := s.api4ConversationDAO.GetBySessionID(ctx, sessionID, canvasID) + session, sessionErr := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, canvasID) if sessionErr != nil { return nil, fmt.Errorf("RunAgent: load session %q: %w: %w", sessionID, sessionErr, ErrAgentStorageError) } @@ -1545,7 +1545,7 @@ func (s *AgentService) createAgentRunSession( if versionRow != nil { session.VersionTitle = versionRow.Title } - return s.api4ConversationDAO.Create(ctx, session) + return s.api4ConversationDAO.Create(ctx, dao.DB, session) } // runIDFor builds the per-run CanvasState identifier: canvasID @@ -1595,7 +1595,7 @@ func (s *AgentService) persistAgentRunSession( if sessionID == "" || s == nil || s.api4ConversationDAO == nil || dao.DB == nil { return nil } - session, err := s.api4ConversationDAO.GetBySessionID(ctx, sessionID, agentID) + session, err := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, agentID) if err != nil { common.Warn("agent run: load session for update failed", zap.String("agent_id", agentID), zap.String("session_id", sessionID), zap.Error(err)) return nil @@ -1622,7 +1622,7 @@ func (s *AgentService) persistAgentRunSession( if state != nil { session.DSL = buildPersistedAgentDSL(runDSL, state) } - return s.api4ConversationDAO.Update(ctx, session) + return s.api4ConversationDAO.Update(ctx, dao.DB, session) } func buildPersistedAgentDSL(runDSL map[string]any, state *canvas.CanvasState) entity.JSONMap { diff --git a/internal/service/agent_sessions.go b/internal/service/agent_sessions.go index 6b029e2f7d..289a75b309 100644 --- a/internal/service/agent_sessions.go +++ b/internal/service/agent_sessions.go @@ -325,18 +325,18 @@ func (s *AgentService) ListAgentSessions(ctx context.Context, userID, tenantID, // shape here instead of a 500 record not found so the // front-end does not log a server error for unknown ids. if errors.Is(err, dao.ErrUserCanvasNotFound) { - return nil, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return nil, common.CodeOperatingError, errors.New("agent not found or no permission") } return nil, common.CodeServerError, fmt.Errorf("failed to check agent permission: %w", err) } if !ok { - return nil, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return nil, common.CodeOperatingError, errors.New("agent not found or no permission") } sessionDAO := dao.NewChatSessionDAO() if req.ExpUserID != "" { - rows, err := sessionDAO.ListAgentSessionNames(ctx, agentID, req.ExpUserID) + rows, err := sessionDAO.ListAgentSessionNames(ctx, dao.DB, agentID, req.ExpUserID) if err != nil { return nil, common.CodeServerError, err } @@ -352,7 +352,7 @@ func (s *AgentService) ListAgentSessions(ctx context.Context, userID, tenantID, return nil, common.CodeArgumentError, err } - total, sessions, err := sessionDAO.ListAgentSessions(ctx, dao.ListAgentSessionsParams{ + total, sessions, err := sessionDAO.ListAgentSessions(ctx, dao.DB, dao.ListAgentSessionsParams{ AgentID: agentID, Page: req.Page, PageSize: req.PageSize, @@ -385,15 +385,15 @@ func (s *AgentService) GetAgentSession(ctx context.Context, userID, agentID, ses ok, err := s.CheckCanvasAccess(userID, agentID) if err != nil { if errors.Is(err, dao.ErrUserCanvasNotFound) { - return nil, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return nil, common.CodeOperatingError, errors.New("agent not found or no permission") } return nil, common.CodeServerError, fmt.Errorf("failed to check agent permission: %w", err) } if !ok { - return nil, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return nil, common.CodeOperatingError, errors.New("agent not found or no permission") } - data, err := s.api4ConversationDAO.GetBySessionID(ctx, sessionID, agentID) + data, err := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, agentID) if err != nil { return nil, common.CodeServerError, fmt.Errorf("failed to fetch session: %w", err) } @@ -411,15 +411,15 @@ func (s *AgentService) DeleteAgentSessionItem(ctx context.Context, userID, agent ok, err := s.CheckCanvasAccess(userID, agentID) if err != nil { if errors.Is(err, dao.ErrUserCanvasNotFound) { - return false, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return false, common.CodeOperatingError, errors.New("agent not found or no permission") } return false, common.CodeServerError, fmt.Errorf("failed to check agent permission: %w", err) } if !ok { - return false, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return false, common.CodeOperatingError, errors.New("agent not found or no permission") } - row, err := s.api4ConversationDAO.DeleteBySessionIDAndAgentID(ctx, sessionID, agentID) + row, err := s.api4ConversationDAO.DeleteBySessionIDAndAgentID(ctx, dao.DB, sessionID, agentID) if err != nil { return false, common.CodeServerError, err } @@ -443,7 +443,7 @@ func (s *AgentService) DeleteAgentSessions(ctx context.Context, userID, agentID // access, which is too permissive for this operation. canvas, err := s.canvasDAO.GetByID(agentID) if err != nil || canvas == nil || canvas.UserID != userID { - return nil, common.CodeDataError, fmt.Errorf("You don't own the agent %s", agentID) + return nil, common.CodeDataError, fmt.Errorf("you don't own the agent %s", agentID) } if len(ids) == 0 { @@ -451,7 +451,7 @@ func (s *AgentService) DeleteAgentSessions(ctx context.Context, userID, agentID return &DeleteAgentSessionsResult{}, common.CodeSuccess, nil } - ids, err = s.api4ConversationDAO.ListIDsByAgentID(ctx, agentID) + ids, err = s.api4ConversationDAO.ListIDsByAgentID(ctx, dao.DB, agentID) if err != nil { return nil, common.CodeServerError, err } @@ -471,7 +471,7 @@ func (s *AgentService) DeleteAgentSessions(ctx context.Context, userID, agentID continue } - conv, err := s.api4ConversationDAO.GetBySessionID(ctx, sessionID, agentID) + conv, err := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, sessionID, agentID) if err != nil { return nil, common.CodeServerError, err } @@ -480,7 +480,7 @@ func (s *AgentService) DeleteAgentSessions(ctx context.Context, userID, agentID continue } - if _, err := s.api4ConversationDAO.DeleteBySessionIDAndAgentID(ctx, sessionID, agentID); err != nil { + if _, err := s.api4ConversationDAO.DeleteBySessionIDAndAgentID(ctx, dao.DB, sessionID, agentID); err != nil { return nil, common.CodeServerError, err } successCount++ @@ -548,7 +548,7 @@ func (s *AgentService) ListAgentTags(userID, canvasCategory string) ([]AgentTagC } // normalizeAgentTags returns an error for unsupported tag payload types. -// The branch behaviour intentionally mirrors the Python implementation: +// The branch behavior intentionally mirrors the Python implementation: // - string: treat the value as a CSV — split on "," and use each piece // as a separate tag ("alpha,beta" → ["alpha", "beta"]). // - []string / []interface{}: the caller already chose the boundary; @@ -627,12 +627,12 @@ func (s *AgentService) UpdateAgentTags(userID, canvasID string, tags interface{} ok, err := s.CheckCanvasAccess(userID, canvasID) if err != nil { if errors.Is(err, dao.ErrUserCanvasNotFound) { - return false, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return false, common.CodeOperatingError, errors.New("agent not found or no permission") } return false, common.CodeServerError, fmt.Errorf("failed to check agent permission: %w", err) } if !ok { - return false, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return false, common.CodeOperatingError, errors.New("agent not found or no permission") } normalized, nErr := normalizeAgentTags(tags) @@ -645,7 +645,7 @@ func (s *AgentService) UpdateAgentTags(userID, canvasID string, tags interface{} } if rows == 0 { if _, getErr := s.canvasDAO.GetByCanvasID(canvasID); getErr != nil { - return false, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return false, common.CodeOperatingError, errors.New("agent not found or no permission") } return true, common.CodeSuccess, nil } @@ -665,12 +665,12 @@ type CreateAgentSessionRequest struct { // CreateAgentSession inserts a fresh conversation row tied to the // given agent canvas. The Phase 5 stub intentionally does NOT run -// Canvas(dsl).reset() (eino runtime is still unimplemented in the Go +// Canvas(dsl).reset() (Eino runtime is still unimplemented in the Go // port); instead it stores a minimal but well-shaped row so that // subsequent ListAgentSessions / GetAgentSession / chat-completion // stubs can return a stable id and the integration suite can verify -// the create + read + delete cycle without depending on a real LLM -// run. When eino lands, the function will gain a pre-run prologue +// the creation + read + delete cycle without depending on a real LLM +// run. When Eino lands, the function will gain a pre-run prologue // pass that calls Canvas.Reset() and stores the assistant message. // // Required columns (per the API4Conversation entity, see @@ -678,7 +678,7 @@ type CreateAgentSessionRequest struct { // - id : 32-hex uuid, matches Python uuid.uuid4().hex // - dialog_id : agent canvas id // - user_id : caller's id -// - message : JSON array (default []); GET path normalises it +// - message : JSON array (default []); GET path normalizes it // - reference : JSON object (default {}) so GET-side parsing // does not crash on .chunks // - dsl : JSON map; copied from user_canvas.dsl if the @@ -703,7 +703,7 @@ func (s *AgentService) CreateAgentSession(ctx context.Context, req *CreateAgentS return nil, common.CodeServerError, fmt.Errorf("check canvas access: %w", err) } if !ok { - return nil, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return nil, common.CodeOperatingError, errors.New("agent not found or no permission") } messages := req.Messages @@ -720,7 +720,7 @@ func (s *AgentService) CreateAgentSession(ctx context.Context, req *CreateAgentS canvas, gErr := s.canvasDAO.GetByID(req.AgentID) if gErr != nil { if errors.Is(gErr, gorm.ErrRecordNotFound) { - return nil, common.CodeOperatingError, errors.New("Agent not found or no permission.") + return nil, common.CodeOperatingError, errors.New("agent not found or no permission") } return nil, common.CodeServerError, fmt.Errorf("load canvas dsl: %w", gErr) } @@ -752,7 +752,7 @@ func (s *AgentService) CreateAgentSession(ctx context.Context, req *CreateAgentS Source: sourcePtr, DSL: dsl, } - if err := s.api4ConversationDAO.Create(ctx, row); err != nil { + if err := s.api4ConversationDAO.Create(ctx, dao.DB, row); err != nil { return nil, common.CodeServerError, fmt.Errorf("create agent session: %w", err) } return row, common.CodeSuccess, nil diff --git a/internal/service/api_token.go b/internal/service/api_token.go index 18ca3f532a..acbc704732 100644 --- a/internal/service/api_token.go +++ b/internal/service/api_token.go @@ -40,7 +40,7 @@ type APIKeyResponse struct { // ListAPIKeys list all API keys for a tenant func (s *SystemService) ListAPIKeys(ctx context.Context, tenantID string) ([]*APIKeyResponse, error) { APITokenDAO := dao.NewAPITokenDAO() - keys, err := APITokenDAO.GetByTenantID(ctx, tenantID) + keys, err := APITokenDAO.GetByTenantID(ctx, dao.DB, tenantID) if err != nil { return nil, err } @@ -98,7 +98,7 @@ func (s *SystemService) CreateAPIKey(ctx context.Context, tenantID string, req * Beta: &betaAPIKey, } - if err := APITokenDAO.Create(ctx, APIKeyData); err != nil { + if err := APITokenDAO.Create(ctx, dao.DB, APIKeyData); err != nil { return nil, err } @@ -117,6 +117,6 @@ func (s *SystemService) CreateAPIKey(ctx context.Context, tenantID string, req * func (s *SystemService) DeleteAPIKey(ctx context.Context, tenantID, key string) error { APITokenDAO := dao.NewAPITokenDAO() - _, err := APITokenDAO.DeleteByTenantIDAndToken(ctx, tenantID, key) + _, err := APITokenDAO.DeleteByTenantIDAndToken(ctx, dao.DB, tenantID, key) return err } diff --git a/internal/service/bot.go b/internal/service/bot.go index adabed4a84..d26a267e30 100644 --- a/internal/service/bot.go +++ b/internal/service/bot.go @@ -83,7 +83,7 @@ func NewBotService(agentSvc *AgentService, llmSvc *LLMService) *BotService { func (s *BotService) ChatbotInfo(ctx context.Context, tenantID, dialogID string) ( title, avatar, prologue, llmID string, hasTavilyKey bool, ec common.ErrorCode, err error, ) { - dialog, err := s.chatDAO.GetDialogByID(ctx, dialogID) + dialog, err := s.chatDAO.GetDialogByID(ctx, dao.DB, dialogID) if err != nil { return "", "", "", "", false, common.CodeDataError, err } diff --git a/internal/service/bot_completion.go b/internal/service/bot_completion.go index 758289f379..93db35c333 100644 --- a/internal/service/bot_completion.go +++ b/internal/service/bot_completion.go @@ -38,6 +38,7 @@ import ( "encoding/json" "errors" "net/http" + "ragflow/internal/dao" "ragflow/internal/utility" "strings" "time" @@ -302,7 +303,7 @@ func (s *BotService) ChatbotCompletion( // ChatSessionDAO.GetDialogByID already filters by status = "1" // so a returned row is valid; we still nil-check defensively // before dereferencing for symmetry with the session path. - dialog, err := s.chatDAO.GetDialogByID(ctx, dialogID) + dialog, err := s.chatDAO.GetDialogByID(ctx, dao.DB, dialogID) if err != nil || dialog == nil || dialog.TenantID != tenantID || dialog.Status == nil || *dialog.Status != common.StatusDialogValid { @@ -351,7 +352,7 @@ func (s *BotService) ChatbotCompletion( UserID: tenantID, Message: seedMsg, } - if err = s.api4ConversationDAO.Create(ctx, session); err != nil { + if err = s.api4ConversationDAO.Create(ctx, dao.DB, session); err != nil { return nil, common.CodeServerError, err } @@ -375,7 +376,7 @@ func (s *BotService) ChatbotCompletion( return out, common.CodeSuccess, nil } - session, err := s.api4ConversationDAO.GetBySessionID(ctx, req.SessionID, dialogID) + session, err := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, req.SessionID, dialogID) if err != nil { return nil, common.CodeServerError, err } @@ -580,7 +581,7 @@ func (s *BotService) persistChatbotTurn( lock := s.persistLock(session.ID) lock.Lock() defer lock.Unlock() - fresh, err := s.api4ConversationDAO.GetBySessionID(ctx, session.ID, session.DialogID) + fresh, err := s.api4ConversationDAO.GetBySessionID(ctx, dao.DB, session.ID, session.DialogID) if err != nil { return err } @@ -632,7 +633,7 @@ func (s *BotService) persistChatbotTurn( } session.Reference = rawRef - return s.api4ConversationDAO.Update(ctx, session) + return s.api4ConversationDAO.Update(ctx, dao.DB, session) } // normalizeBotBoolFlag coerces the JSON-encoded reasoning / internet diff --git a/internal/service/bot_completion_history_test.go b/internal/service/bot_completion_history_test.go index 684d8d46cf..2d41b64713 100644 --- a/internal/service/bot_completion_history_test.go +++ b/internal/service/bot_completion_history_test.go @@ -408,7 +408,7 @@ func TestBotService_ChatbotCompletion_NewSessionSkipsLLM(t *testing.T) { // The persisted session must hold exactly the prologue turn — // no empty user message may be written. ctx := t.Context() - row, derr := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, got[0].SessionID, "dlg-1") + row, derr := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, got[0].SessionID, "dlg-1") if derr != nil || row == nil { t.Fatalf("session not persisted: row=%v err=%v", row, derr) } @@ -473,7 +473,7 @@ func TestPersistChatbotTurn_AppendsPairAndReference(t *testing.T) { Message: seed, } ctx := t.Context() - if err := dao.NewAPI4ConversationDAO().Create(ctx, sess); err != nil { + if err := dao.NewAPI4ConversationDAO().Create(ctx, dao.DB, sess); err != nil { t.Fatalf("seed session: %v", err) } @@ -486,7 +486,7 @@ func TestPersistChatbotTurn_AppendsPairAndReference(t *testing.T) { t.Fatalf("persistChatbotTurn: %v", err) } - row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, "sess-p1", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p1", "dlg-p1") if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } @@ -525,7 +525,7 @@ func TestPersistChatbotTurn_NilReferenceDefaultsToEmpty(t *testing.T) { ctx := t.Context() sess := &entity.API4Conversation{ID: "sess-p2", DialogID: "dlg-p1", UserID: "tenant-1"} - if err := dao.NewAPI4ConversationDAO().Create(ctx, sess); err != nil { + if err := dao.NewAPI4ConversationDAO().Create(ctx, dao.DB, sess); err != nil { t.Fatalf("seed session: %v", err) } @@ -534,7 +534,7 @@ func TestPersistChatbotTurn_NilReferenceDefaultsToEmpty(t *testing.T) { t.Fatalf("persistChatbotTurn: %v", err) } - row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, "sess-p2", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p2", "dlg-p1") if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } @@ -571,7 +571,7 @@ func TestPersistChatbotTurn_ConcurrentSameSession(t *testing.T) { ctx := t.Context() sess := &entity.API4Conversation{ID: "sess-p3", DialogID: "dlg-p1", UserID: "tenant-1"} - if err = dao.NewAPI4ConversationDAO().Create(ctx, sess); err != nil { + if err = dao.NewAPI4ConversationDAO().Create(ctx, dao.DB, sess); err != nil { t.Fatalf("seed session: %v", err) } @@ -590,7 +590,7 @@ func TestPersistChatbotTurn_ConcurrentSameSession(t *testing.T) { } wg.Wait() - row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, "sess-p3", "dlg-p1") + row, err := dao.NewAPI4ConversationDAO().GetBySessionID(ctx, dao.DB, "sess-p3", "dlg-p1") if err != nil || row == nil { t.Fatalf("re-read session: row=%v err=%v", row, err) } diff --git a/internal/service/chat.go b/internal/service/chat.go index d479ee2f61..939ed6f94e 100644 --- a/internal/service/chat.go +++ b/internal/service/chat.go @@ -87,6 +87,7 @@ func (s *ChatService) ListChats(ctx context.Context, userID, status, keywords st if len(ownerIDs) == 0 { chats, total, err = s.chatDAO.ListByTenantIDs( ctx, + dao.DB, nil, userID, page, @@ -111,7 +112,7 @@ func (s *ChatService) ListChats(ctx context.Context, userID, status, keywords st }, nil } - chats, total, err = s.chatDAO.ListByOwnerIDs(ctx, filterOwnerIDs, userID, orderBy, desc, keywords) + chats, total, err = s.chatDAO.ListByOwnerIDs(ctx, dao.DB, filterOwnerIDs, userID, orderBy, desc, keywords) if err != nil { return nil, err } @@ -213,7 +214,7 @@ func (s *ChatService) Create(ctx context.Context, userID string, req map[string] } if tenantValue, ok := req["tenant_id"]; ok && isTruthy(tenantValue) { - return nil, common.CodeDataError, errors.New("`tenant_id` must not be provided.") + return nil, common.CodeDataError, errors.New("`tenant_id` must not be provided") } name, err := validateCreateChatName(req["name"]) @@ -256,13 +257,13 @@ func (s *ChatService) Create(ctx context.Context, userID string, req map[string] if promptConfigValue, ok := req["prompt_config"]; ok { if _, ok := mapFromValue(promptConfigValue); !ok { - return nil, common.CodeDataError, errors.New("`prompt_config` should be an object.") + return nil, common.CodeDataError, errors.New("`prompt_config` should be an object") } } if metaDataFilterValue, ok := req["meta_data_filter"]; ok && metaDataFilterValue != nil { if _, ok := mapFromValue(metaDataFilterValue); !ok { - return nil, common.CodeDataError, errors.New("`meta_data_filter` should be an object.") + return nil, common.CodeDataError, errors.New("`meta_data_filter` should be an object") } } @@ -320,20 +321,20 @@ func (s *ChatService) Create(ctx context.Context, userID string, req map[string] applyCreatePromptDefaults(req) filterCreateChatPersistedFields(req) - exists, err := s.chatDAO.ExistsByNameTenantStatus(ctx, name, userID, string(entity.StatusValid)) + exists, err := s.chatDAO.ExistsByNameTenantStatus(ctx, dao.DB, name, userID, string(entity.StatusValid)) if err != nil { return nil, common.CodeServerError, err } if exists { - return nil, common.CodeDataError, errors.New("Duplicated chat name in creating chat.") + return nil, common.CodeDataError, errors.New("duplicated chat name in creating chat") } chat := buildCreateChatEntity(req, userID) - if err = s.chatDAO.Create(ctx, chat); err != nil { + if err = s.chatDAO.Create(ctx, dao.DB, chat); err != nil { return nil, common.CodeDataError, errors.New("failed to create chat") } - chat, err = s.chatDAO.GetByID(ctx, chat.ID) + chat, err = s.chatDAO.GetByID(ctx, dao.DB, chat.ID) if err != nil { return nil, common.CodeDataError, errors.New("failed to retrieve created chat") } @@ -347,18 +348,18 @@ func (s *ChatService) Create(ctx context.Context, userID string, req map[string] func validateCreateChatName(value interface{}) (string, error) { if value == nil { - return "", errors.New("`name` is required.") + return "", errors.New("`name` is required") } name, ok := value.(string) if !ok { - return "", errors.New("Chat name must be a string.") + return "", errors.New("chat name must be a string") } name = strings.TrimSpace(name) if name == "" { - return "", errors.New("`name` is required.") + return "", errors.New("`name` is required") } if len([]byte(name)) > 255 { - return "", fmt.Errorf("Chat name length is %d which is larger than 255.", len([]byte(name))) + return "", fmt.Errorf("chat name length is %d which is larger than 255", len([]byte(name))) } return name, nil } @@ -785,7 +786,7 @@ func (s *ChatService) splitModelNameAndFactory(embeddingModelID string) string { } func (s *ChatService) getOwnedValidChat(ctx context.Context, userID, chatID string) (*entity.Chat, error) { - chat, err := s.chatDAO.GetByIDAndStatus(ctx, chatID, string(entity.StatusValid)) + chat, err := s.chatDAO.GetByIDAndStatus(ctx, dao.DB, chatID, string(entity.StatusValid)) if err != nil { return nil, errors.New("no authorization") } @@ -852,7 +853,7 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, } if !patch && isTruthy(req["tenant_id"]) { - return nil, errors.New("`tenant_id` must not be provided.") + return nil, errors.New("`tenant_id` must not be provided") } if value, ok := req["name"]; ok { @@ -912,7 +913,7 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, if value, ok := req["prompt_config"]; ok { promptConfig, ok := mapFromValue(value) if !ok { - return nil, errors.New("`prompt_config` should be an object.") + return nil, errors.New("`prompt_config` should be an object") } if patch { req["prompt_config"] = mergeJSONMap(currentChat.PromptConfig, promptConfig) @@ -935,7 +936,7 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, } else { metaDataFilter, ok := mapFromValue(value) if !ok { - return nil, errors.New("`meta_data_filter` should be an object.") + return nil, errors.New("`meta_data_filter` should be an object") } req["meta_data_filter"] = entity.JSONMap(metaDataFilter) } @@ -951,20 +952,20 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, currentName = *currentChat.Name } if strings.ToLower(name) != strings.ToLower(currentName) { - existingNames, err := s.chatDAO.GetExistingNames(ctx, userID, string(entity.StatusValid)) + existingNames, err := s.chatDAO.GetExistingNames(ctx, dao.DB, userID, string(entity.StatusValid)) if err != nil { return nil, err } for _, existingName := range existingNames { if strings.EqualFold(existingName, name) { - return nil, errors.New("Duplicated chat name.") + return nil, errors.New("duplicated chat name") } } } } if len(updates) > 0 { - if err = s.chatDAO.UpdateByID(ctx, chatID, updates); err != nil { + if err = s.chatDAO.UpdateByID(ctx, dao.DB, chatID, updates); err != nil { if patch { return nil, errors.New("failed to update chat") } @@ -972,7 +973,7 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, } } - updatedChat, err := s.chatDAO.GetByID(ctx, chatID) + updatedChat, err := s.chatDAO.GetByID(ctx, dao.DB, chatID) if err != nil { return nil, errors.New("failed to retrieve updated chat") } @@ -982,23 +983,23 @@ func (s *ChatService) updateChatREST(ctx context.Context, userID, chatID string, func validateRESTChatName(value interface{}, required bool) (string, bool, error) { if value == nil { if required { - return "", false, errors.New("`name` is required.") + return "", false, errors.New("`name` is required") } return "", false, nil } name, ok := value.(string) if !ok { - return "", false, errors.New("Chat name must be a string.") + return "", false, errors.New("chat name must be a string") } name = strings.TrimSpace(name) if name == "" { if required { - return "", false, errors.New("`name` is required.") + return "", false, errors.New("`name` is required") } - return "", false, errors.New("`name` cannot be empty.") + return "", false, errors.New("`name` cannot be empty") } if len([]byte(name)) > 255 { - return "", false, fmt.Errorf("Chat name length is %d which is larger than 255.", len([]byte(name))) + return "", false, fmt.Errorf("chat name length is %d which is larger than 255", len([]byte(name))) } return name, true, nil } @@ -1157,7 +1158,7 @@ func (s *ChatService) DeleteChat(ctx context.Context, userID, chatID string) err if _, err := s.getOwnedValidChat(ctx, userID, chatID); err != nil { return err } - if err := s.chatDAO.UpdateByID(ctx, chatID, map[string]interface{}{ + if err := s.chatDAO.UpdateByID(ctx, dao.DB, chatID, map[string]interface{}{ "status": string(entity.StatusInvalid), }); err != nil { return fmt.Errorf("failed to delete chat %s", chatID) @@ -1201,7 +1202,7 @@ func checkDuplicateChatIDs(ids []string) ([]string, []string) { func (s *ChatService) BulkDeleteChats(ctx context.Context, userID string, req *BulkDeleteChatsRequest) (map[string]interface{}, error) { ids := req.IDs if len(ids) == 0 && req.DeleteAll { - chats, err := s.chatDAO.ListByTenantID(ctx, userID, string(entity.StatusValid)) + chats, err := s.chatDAO.ListByTenantID(ctx, dao.DB, userID, string(entity.StatusValid)) if err != nil { return nil, err } @@ -1223,7 +1224,7 @@ func (s *ChatService) BulkDeleteChats(ctx context.Context, userID string, req *B errorsList = append(errorsList, fmt.Sprintf("Chat(%s) not found.", chatID)) continue } - if err := s.chatDAO.UpdateByID(ctx, chatID, map[string]interface{}{ + if err := s.chatDAO.UpdateByID(ctx, dao.DB, chatID, map[string]interface{}{ "status": string(entity.StatusInvalid), }); err != nil { errorsList = append(errorsList, fmt.Sprintf("Failed to delete chat %s", chatID)) @@ -1275,7 +1276,7 @@ func (s *ChatService) GetChat(ctx context.Context, userID string, chatID string) // Python: for tenant in tenants: if DialogService.query(tenant_id=tenant.tenant_id, id=chat_id, status=StatusEnum.VALID.value): break hasPermission := false for _, tenant := range tenants { - chats, err := s.chatDAO.QueryByTenantIDAndID(ctx, tenant.TenantID, chatID, "1") + chats, err := s.chatDAO.QueryByTenantIDAndID(ctx, dao.DB, tenant.TenantID, chatID, "1") if err != nil { continue // Try next tenant } @@ -1290,7 +1291,7 @@ func (s *ChatService) GetChat(ctx context.Context, userID string, chatID string) } // Step 3: Get chat detail (same as Python DialogService.get_by_id(chat_id)) - chat, err := s.chatDAO.GetByID(ctx, chatID) + chat, err := s.chatDAO.GetByID(ctx, dao.DB, chatID) if err != nil { return nil, fmt.Errorf("chat not found") } diff --git a/internal/service/chat_channel.go b/internal/service/chat_channel.go index b6b2773775..63c4af526e 100644 --- a/internal/service/chat_channel.go +++ b/internal/service/chat_channel.go @@ -49,31 +49,31 @@ func (s *ChatChannelService) Insert(ctx context.Context, channel *entity.ChatCha if channel.Status == 0 { channel.Status = 1 } - return s.chatChannelDAO.Create(ctx, channel) + return s.chatChannelDAO.Create(ctx, dao.DB, channel) } func (s *ChatChannelService) GetByID(ctx context.Context, id string) (*entity.ChatChannel, error) { if id == "" { return nil, errors.New("id is empty") } - return s.chatChannelDAO.GetByIDOnly(ctx, id) + return s.chatChannelDAO.GetByIDOnly(ctx, dao.DB, id) } func (s *ChatChannelService) List(ctx context.Context, tenantID string) ([]*entity.ChatChannelListResponse, error) { - return s.chatChannelDAO.ListByTenantID(ctx, tenantID) + return s.chatChannelDAO.ListByTenantID(ctx, dao.DB, tenantID) } func (s *ChatChannelService) CreateChatChannel(ctx context.Context, tenantID, name, channelType string, config entity.JSONMap, chatID *string) (*entity.ChatChannel, error) { if chatID != nil && *chatID != "" { - dialog, err := s.chatDAO.GetByID(ctx, *chatID) + dialog, err := s.chatDAO.GetByID(ctx, dao.DB, *chatID) if err != nil { if dao.IsNotFoundErr(err) { - return nil, errors.New("Can't find this chat assistant!") + return nil, errors.New("can't find this chat assistant") } return nil, err } if dialog.TenantID != tenantID { - return nil, errors.New("No authorization.") + return nil, errors.New("no authorization") } } row := &entity.ChatChannel{ @@ -98,7 +98,7 @@ func (s *ChatChannelService) CreateChatChannel(ctx context.Context, tenantID, na } func (s *ChatChannelService) accessible(ctx context.Context, userID, channelID string) (*entity.ChatChannel, bool, error) { - channel, err := s.chatChannelDAO.GetByIDOnly(ctx, channelID) + channel, err := s.chatChannelDAO.GetByIDOnly(ctx, dao.DB, channelID) if err != nil { if dao.IsNotFoundErr(err) { return nil, false, nil @@ -129,13 +129,13 @@ func (s *ChatChannelService) GetChatChannel(ctx context.Context, userID, channel return nil, common.CodeServerError, err } if !ok { - return nil, common.CodeAuthenticationError, errors.New("No authorization.") + return nil, common.CodeAuthenticationError, errors.New("no authorization") } - channel, err := s.chatChannelDAO.GetByIDOnly(ctx, channelID) + channel, err := s.chatChannelDAO.GetByIDOnly(ctx, dao.DB, channelID) if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("Can't find this chat channel!") + return nil, common.CodeDataError, errors.New("can't find this chat channel") } return nil, common.CodeServerError, err } @@ -148,10 +148,10 @@ func (s *ChatChannelService) UpdateChatChannel(ctx context.Context, userID, chan return nil, common.CodeServerError, err } if !ok { - return nil, common.CodeAuthenticationError, errors.New("No authorization.") + return nil, common.CodeAuthenticationError, errors.New("no authorization") } if channel == nil { - return nil, common.CodeDataError, errors.New("Can't find this chat channel!") + return nil, common.CodeDataError, errors.New("can't find this chat channel") } updates := map[string]interface{}{} @@ -185,15 +185,15 @@ func (s *ChatChannelService) UpdateChatChannel(ctx context.Context, userID, chan return nil, common.CodeDataError, errors.New("chat_id must be string or null") } if chatID != "" { - dialog, err := s.chatDAO.GetByID(ctx, chatID) + dialog, err := s.chatDAO.GetByID(ctx, dao.DB, chatID) if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("Can't find this chat assistant!") + return nil, common.CodeDataError, errors.New("can't find this chat assistant") } return nil, common.CodeServerError, err } if dialog.TenantID != channel.TenantID { - return nil, common.CodeAuthenticationError, errors.New("No authorization.") + return nil, common.CodeAuthenticationError, errors.New("no authorization") } } updates["chat_id"] = chatID @@ -201,15 +201,15 @@ func (s *ChatChannelService) UpdateChatChannel(ctx context.Context, userID, chan } if len(updates) > 0 { - if err := s.chatChannelDAO.UpdateByID(ctx, channelID, channel.TenantID, updates); err != nil { + if err = s.chatChannelDAO.UpdateByID(ctx, dao.DB, channelID, channel.TenantID, updates); err != nil { return nil, common.CodeDataError, err } } - updated, err := s.chatChannelDAO.GetByIDOnly(ctx, channelID) + updated, err := s.chatChannelDAO.GetByIDOnly(ctx, dao.DB, channelID) if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("Can't find this chat channel!") + return nil, common.CodeDataError, errors.New("can't find this chat channel") } return nil, common.CodeServerError, err } @@ -222,13 +222,13 @@ func (s *ChatChannelService) DeleteChatChannel(ctx context.Context, userID, chan return false, common.CodeServerError, err } if !ok { - return false, common.CodeAuthenticationError, errors.New("No authorization.") + return false, common.CodeAuthenticationError, errors.New("no authorization") } if channel == nil { - return false, common.CodeAuthenticationError, errors.New("No authorization.") + return false, common.CodeAuthenticationError, errors.New("no authorization") } - if err = s.chatChannelDAO.DeleteByID(ctx, channelID, channel.TenantID); err != nil { + if err = s.chatChannelDAO.DeleteByID(ctx, dao.DB, channelID, channel.TenantID); err != nil { return false, common.CodeDataError, err } return true, common.CodeSuccess, nil diff --git a/internal/service/chat_channel_test.go b/internal/service/chat_channel_test.go index 1117fa7867..5cedcd7509 100644 --- a/internal/service/chat_channel_test.go +++ b/internal/service/chat_channel_test.go @@ -200,7 +200,7 @@ func TestChatChannelServiceUpdateChatChannelRejectsCrossTenantDialog(t *testing. if code != common.CodeAuthenticationError { t.Fatalf("expected authentication error, got %v", code) } - if err == nil || !strings.Contains(err.Error(), "No authorization.") { + if err == nil || !strings.Contains(err.Error(), "no authorization") { t.Fatalf("expected authorization error, got %v", err) } } @@ -223,7 +223,7 @@ func TestChatChannelServiceDeleteChatChannel(t *testing.T) { if code != common.CodeAuthenticationError { t.Fatalf("expected authentication error, got %v", code) } - if err == nil || !strings.Contains(err.Error(), "No authorization.") { + if err == nil || !strings.Contains(err.Error(), "no authorization") { t.Fatalf("expected authorization error, got %v", err) } if deleted { diff --git a/internal/service/chat_delete_test.go b/internal/service/chat_delete_test.go index 6b6a8dbac9..a7508b0cb9 100644 --- a/internal/service/chat_delete_test.go +++ b/internal/service/chat_delete_test.go @@ -66,7 +66,7 @@ func TestChatServiceDeleteChatRejectsNonOwner(t *testing.T) { t.Fatalf("expected no authorization, got %v", err) } - chat, getErr := svc.chatDAO.GetByID(ctx, "chat-1") + chat, getErr := svc.chatDAO.GetByID(ctx, dao.DB, "chat-1") if getErr != nil { t.Fatalf("failed to fetch chat: %v", getErr) } @@ -92,7 +92,7 @@ func TestChatServiceBulkDeleteChatsDeleteAllOnlyDeletesOwnedChats(t *testing.T) t.Fatalf("expected success_count 2, got %+v", result["success_count"]) } - owned1, err := svc.chatDAO.GetByID(ctx, "chat-1") + owned1, err := svc.chatDAO.GetByID(ctx, dao.DB, "chat-1") if err != nil { t.Fatalf("failed to fetch chat-1: %v", err) } @@ -100,7 +100,7 @@ func TestChatServiceBulkDeleteChatsDeleteAllOnlyDeletesOwnedChats(t *testing.T) t.Fatalf("expected chat-1 invalid, got %+v", owned1.Status) } - owned2, err := svc.chatDAO.GetByID(ctx, "chat-2") + owned2, err := svc.chatDAO.GetByID(ctx, dao.DB, "chat-2") if err != nil { t.Fatalf("failed to fetch chat-2: %v", err) } @@ -108,7 +108,7 @@ func TestChatServiceBulkDeleteChatsDeleteAllOnlyDeletesOwnedChats(t *testing.T) t.Fatalf("expected chat-2 invalid, got %+v", owned2.Status) } - other, err := svc.chatDAO.GetByID(ctx, "chat-3") + other, err := svc.chatDAO.GetByID(ctx, dao.DB, "chat-3") if err != nil { t.Fatalf("failed to fetch chat-3: %v", err) } diff --git a/internal/service/chat_rest_update_test.go b/internal/service/chat_rest_update_test.go index 6d8c9693a4..023153dbf9 100644 --- a/internal/service/chat_rest_update_test.go +++ b/internal/service/chat_rest_update_test.go @@ -141,7 +141,7 @@ func TestChatServiceCreateDefaultsMetaDataFilter(t *testing.T) { if !ok || chatID == "" { t.Fatalf("expected created chat id, got %+v", resp["id"]) } - chat, err := svc.chatDAO.GetByID(ctx, chatID) + chat, err := svc.chatDAO.GetByID(ctx, dao.DB, chatID) if err != nil { t.Fatalf("failed to fetch created chat: %v", err) } @@ -178,7 +178,7 @@ func TestChatServiceCreateRejectsInvalidMetaDataFilter(t *testing.T) { "name": "created chat", "meta_data_filter": []interface{}{"invalid"}, }) - if err == nil || err.Error() != "`meta_data_filter` should be an object." { + if err == nil || err.Error() != "`meta_data_filter` should be an object" { t.Fatalf("expected meta_data_filter error, got %v", err) } if code != common.CodeDataError { @@ -206,7 +206,7 @@ func TestChatServicePatchChatMergesPromptConfigAndLLMSetting(t *testing.T) { t.Fatalf("response should expose dataset_ids: %+v", resp) } - chat, err := svc.chatDAO.GetByID(ctx, "chat-1") + chat, err := svc.chatDAO.GetByID(ctx, dao.DB, "chat-1") if err != nil { t.Fatalf("failed to fetch updated chat: %v", err) } @@ -233,7 +233,7 @@ func TestChatServiceUpdateChatRejectsTenantID(t *testing.T) { _, err := svc.UpdateChat(ctx, "user-1", "chat-1", map[string]interface{}{ "tenant_id": "tenant-2", }) - if err == nil || err.Error() != "`tenant_id` must not be provided." { + if err == nil || err.Error() != "`tenant_id` must not be provided" { t.Fatalf("expected tenant_id error, got %v", err) } } @@ -251,7 +251,7 @@ func TestChatServiceUpdateChatRejectsInvalidLLMSetting(t *testing.T) { t.Fatalf("expected llm_setting error, got %v", err) } - chat, err := svc.chatDAO.GetByID(ctx, "chat-1") + chat, err := svc.chatDAO.GetByID(ctx, dao.DB, "chat-1") if err != nil { t.Fatalf("failed to fetch chat: %v", err) } @@ -277,7 +277,7 @@ func TestChatServiceUpdateChatAcceptsMetaDataFilterObject(t *testing.T) { t.Fatalf("UpdateChat failed: %v", err) } - chat, err := svc.chatDAO.GetByID(ctx, "chat-1") + chat, err := svc.chatDAO.GetByID(ctx, dao.DB, "chat-1") if err != nil { t.Fatalf("failed to fetch chat: %v", err) } @@ -300,7 +300,7 @@ func TestChatServiceUpdateChatBackfillsNilMetaDataFilter(t *testing.T) { } assertEmptyMetaDataFilter(t, resp["meta_data_filter"]) - chat, err := svc.chatDAO.GetByID(ctx, "chat-1") + chat, err := svc.chatDAO.GetByID(ctx, dao.DB, "chat-1") if err != nil { t.Fatalf("failed to fetch chat: %v", err) } @@ -324,7 +324,7 @@ func TestChatServicePatchChatIgnoresTenantIDAndUpdatesName(t *testing.T) { t.Fatalf("PatchChat failed: %v", err) } - chat, err := svc.chatDAO.GetByID(ctx, "chat-1") + chat, err := svc.chatDAO.GetByID(ctx, dao.DB, "chat-1") if err != nil { t.Fatalf("failed to fetch updated chat: %v", err) } @@ -348,7 +348,7 @@ func TestChatServiceCreateValidatesName(t *testing.T) { if code != common.CodeDataError { t.Fatalf("expected data error code, got %d", code) } - if err.Error() != "`name` is required." { + if err.Error() != "`name` is required" { t.Fatalf("unexpected error: %v", err) } } @@ -366,7 +366,7 @@ func TestChatServiceCreateRejectsDuplicateName(t *testing.T) { if code != common.CodeDataError { t.Fatalf("expected data error code, got %d", code) } - if !strings.Contains(err.Error(), "Duplicated chat name") { + if !strings.Contains(err.Error(), "duplicated chat name") { t.Fatalf("unexpected error: %v", err) } } @@ -381,7 +381,7 @@ func TestChatServiceUpdateValidatesName(t *testing.T) { if err == nil { t.Fatal("expected name validation error") } - if err.Error() != "`name` is required." { + if err.Error() != "`name` is required" { t.Fatalf("unexpected error: %v", err) } } @@ -397,7 +397,7 @@ func TestChatServiceUpdateRejectsDuplicateName(t *testing.T) { if err == nil { t.Fatal("expected duplicate name error") } - if err.Error() != "Duplicated chat name." { + if err.Error() != "duplicated chat name" { t.Fatalf("unexpected error: %v", err) } } @@ -409,7 +409,7 @@ func TestChatServicePatchChatRejectsEmptyName(t *testing.T) { svc := NewChatService() ctx := t.Context() _, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": ""}) - if err == nil || err.Error() != "`name` cannot be empty." { + if err == nil || err.Error() != "`name` cannot be empty" { t.Fatalf("expected empty name error, got %v", err) } } @@ -421,7 +421,7 @@ func TestChatServicePatchChatRejectsNonStringName(t *testing.T) { svc := NewChatService() ctx := t.Context() _, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": 123}) - if err == nil || err.Error() != "Chat name must be a string." { + if err == nil || err.Error() != "chat name must be a string" { t.Fatalf("expected non-string name error, got %v", err) } } @@ -437,7 +437,7 @@ func TestChatServicePatchChatRejectsTooLongName(t *testing.T) { if err == nil { t.Fatal("expected too long name error") } - if err.Error() != "Chat name length is 256 which is larger than 255." { + if err.Error() != "chat name length is 256 which is larger than 255" { t.Fatalf("unexpected error: %v", err) } } @@ -450,7 +450,7 @@ func TestChatServiceUpdateRejectsDuplicateNameCaseInsensitive(t *testing.T) { svc := NewChatService() ctx := t.Context() _, err := svc.PatchChat(ctx, "user-1", "chat-1", map[string]interface{}{"name": "CHAT-CHAT-2"}) - if err == nil || err.Error() != "Duplicated chat name." { + if err == nil || err.Error() != "duplicated chat name" { t.Fatalf("expected case-insensitive duplicate name error, got %v", err) } } @@ -464,7 +464,7 @@ func TestChatServiceCreateRejectsTenantID(t *testing.T) { "name": "valid chat", "tenant_id": "other-tenant", }) - if err == nil || err.Error() != "`tenant_id` must not be provided." { + if err == nil || err.Error() != "`tenant_id` must not be provided" { t.Fatalf("expected tenant_id error, got %v", err) } if code != common.CodeDataError { @@ -481,7 +481,7 @@ func TestChatServiceCreateRejectsInvalidPromptConfig(t *testing.T) { "name": "valid chat", "prompt_config": "invalid", }) - if err == nil || err.Error() != "`prompt_config` should be an object." { + if err == nil || err.Error() != "`prompt_config` should be an object" { t.Fatalf("expected prompt_config error, got %v", err) } if code != common.CodeDataError { diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go index 3dfc363c02..244cee331a 100644 --- a/internal/service/chat_session.go +++ b/internal/service/chat_session.go @@ -41,14 +41,14 @@ import ( // Interfaces for testability — satisfied by the concrete DAO/pipeline types. type chatSessionStore interface { - GetByID(ctx context.Context, id string) (*entity.ChatSession, error) - GetBySessionIDAndChatID(ctx context.Context, sessionID, chatID string) (*entity.ChatSession, error) - Create(ctx context.Context, conv *entity.ChatSession) error - UpdateByID(ctx context.Context, id string, updates map[string]interface{}) error - DeleteByID(ctx context.Context, id string) error - ListByChatID(ctx context.Context, chatID string) ([]*entity.ChatSession, error) - GetDialogByID(ctx context.Context, chatID string) (*entity.Chat, error) - CheckDialogExists(ctx context.Context, tenantID, chatID string) (bool, error) + GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.ChatSession, error) + GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string) (*entity.ChatSession, error) + Create(ctx context.Context, db *gorm.DB, conv *entity.ChatSession) error + UpdateByID(ctx context.Context, db *gorm.DB, id string, updates map[string]interface{}) error + DeleteByID(ctx context.Context, db *gorm.DB, id string) error + ListByChatID(ctx context.Context, db *gorm.DB, chatID string) ([]*entity.ChatSession, error) + GetDialogByID(ctx context.Context, db *gorm.DB, chatID string) (*entity.Chat, error) + CheckDialogExists(ctx context.Context, db *gorm.DB, tenantID, chatID string) (bool, error) } type userTenantStore interface { @@ -130,17 +130,17 @@ func (s *ChatSessionService) SetChatSession(ctx context.Context, userID string, "name": name, "user_id": userID, } - if err := s.chatSessionDAO.UpdateByID(ctx, req.SessionID, updates); err != nil { + if err := s.chatSessionDAO.UpdateByID(ctx, dao.DB, req.SessionID, updates); err != nil { return nil, errors.New("chat session not found") } - session, err := s.chatSessionDAO.GetByID(ctx, req.SessionID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, req.SessionID) if err != nil { return nil, errors.New("fail to update a chat session") } return &SetChatSessionResponse{ChatSession: session}, nil } - dialog, err := s.chatSessionDAO.GetDialogByID(ctx, req.DialogID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, dao.DB, req.DialogID) if err != nil { return nil, errors.New("dialog not found") } @@ -167,7 +167,7 @@ func (s *ChatSessionService) SetChatSession(ctx context.Context, userID string, UserID: &userID, Reference: referenceJSON, } - if err = s.chatSessionDAO.Create(ctx, session); err != nil { + if err = s.chatSessionDAO.Create(ctx, dao.DB, session); err != nil { return nil, errors.New("fail to create a chat session") } @@ -189,14 +189,14 @@ func (s *ChatSessionService) RemoveChatSessions(ctx context.Context, userID stri tenantIDSet[userID] = true for _, convID := range chatSessions { - session, err := s.chatSessionDAO.GetByID(ctx, convID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, convID) if err != nil { return fmt.Errorf("Chat session not found: %s", convID) } isOwner := false for tenantID := range tenantIDSet { - exists, err := s.chatSessionDAO.CheckDialogExists(ctx, tenantID, session.DialogID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, dao.DB, tenantID, session.DialogID) if err != nil { return err } @@ -209,7 +209,7 @@ func (s *ChatSessionService) RemoveChatSessions(ctx context.Context, userID stri return errors.New("Only owner of chat session authorized for this operation") } - if err = s.chatSessionDAO.DeleteByID(ctx, convID); err != nil { + if err = s.chatSessionDAO.DeleteByID(ctx, dao.DB, convID); err != nil { return err } } @@ -253,7 +253,7 @@ func (s *ChatSessionService) ListChatSessions(ctx context.Context, userID string isOwner := false for _, tenantID := range tenantIDs { var exists bool - exists, err = s.chatSessionDAO.CheckDialogExists(ctx, tenantID, chatID) + exists, err = s.chatSessionDAO.CheckDialogExists(ctx, dao.DB, tenantID, chatID) if err != nil { return nil, err } @@ -266,7 +266,7 @@ func (s *ChatSessionService) ListChatSessions(ctx context.Context, userID string // Also check with userID as tenant if !isOwner { var exists bool - exists, err = s.chatSessionDAO.CheckDialogExists(ctx, userID, chatID) + exists, err = s.chatSessionDAO.CheckDialogExists(ctx, dao.DB, userID, chatID) if err != nil { return nil, err } @@ -278,7 +278,7 @@ func (s *ChatSessionService) ListChatSessions(ctx context.Context, userID string } // List chat sessions - sessions, err := s.chatSessionDAO.ListByChatID(ctx, chatID) + sessions, err := s.chatSessionDAO.ListByChatID(ctx, dao.DB, chatID) if err != nil { return nil, err } @@ -296,7 +296,7 @@ func (s *ChatSessionService) GetSession(ctx context.Context, userID, chatID, ses return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - session, err := s.chatSessionDAO.GetByID(ctx, sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) if err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Session not found!") @@ -307,7 +307,7 @@ func (s *ChatSessionService) GetSession(ctx context.Context, userID, chatID, ses return nil, common.CodeDataError, errors.New("Session does not belong to this chat!") } - dialog, err := s.chatSessionDAO.GetDialogByID(ctx, chatID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, dao.DB, chatID) if err != nil && !isChatSessionNotFound(err) { return nil, common.CodeServerError, err } @@ -325,7 +325,7 @@ func (s *ChatSessionService) CreateSession(ctx context.Context, userID, chatID s return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - dialog, err := s.chatSessionDAO.GetDialogByID(ctx, chatID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, dao.DB, chatID) if err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Chat not found!") @@ -370,11 +370,11 @@ func (s *ChatSessionService) CreateSession(ctx context.Context, userID, chatID s Reference: referenceJSON, } - if err = s.chatSessionDAO.Create(ctx, conv); err != nil { + if err = s.chatSessionDAO.Create(ctx, dao.DB, conv); err != nil { return nil, common.CodeDataError, errors.New("Fail to create a session!") } - session, err := s.chatSessionDAO.GetByID(ctx, conv.ID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, conv.ID) if err != nil { return nil, common.CodeDataError, errors.New("Fail to create a session!") } @@ -399,7 +399,7 @@ func (s *ChatSessionService) DeleteSessions(ctx context.Context, userID, chatID if !hasIDs || len(sessionIDs) == 0 { deleteAll, _ := req["delete_all"].(bool) if deleteAll { - sessions, err := s.chatSessionDAO.ListByChatID(ctx, chatID) + sessions, err := s.chatSessionDAO.ListByChatID(ctx, dao.DB, chatID) if err != nil { return nil, "", common.CodeServerError, err } @@ -420,7 +420,7 @@ func (s *ChatSessionService) DeleteSessions(ctx context.Context, userID, chatID successCount := 0 for _, sid := range uniqueIDs { - session, err := s.chatSessionDAO.GetBySessionIDAndChatID(ctx, sid, chatID) + session, err := s.chatSessionDAO.GetBySessionIDAndChatID(ctx, dao.DB, sid, chatID) if err != nil { errorsList = append(errorsList, fmt.Sprintf("The chat doesn't own the session %s", sid)) continue @@ -428,7 +428,7 @@ func (s *ChatSessionService) DeleteSessions(ctx context.Context, userID, chatID s.removeSessionUploadFiles(userID, session) - if err = s.chatSessionDAO.DeleteByID(ctx, sid); err != nil { + if err = s.chatSessionDAO.DeleteByID(ctx, dao.DB, sid); err != nil { return nil, "", common.CodeServerError, err } @@ -552,7 +552,7 @@ func (s *ChatSessionService) UpdateSession(ctx context.Context, userID, chatID, return nil, common.CodeArgumentError, errors.New("Request body cannot be empty") } - if _, err := s.chatSessionDAO.GetBySessionIDAndChatID(ctx, sessionID, chatID); err != nil { + if _, err := s.chatSessionDAO.GetBySessionIDAndChatID(ctx, dao.DB, sessionID, chatID); err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Session not found!") } @@ -591,14 +591,14 @@ func (s *ChatSessionService) UpdateSession(ctx context.Context, userID, chatID, } } - if err = s.chatSessionDAO.UpdateByID(ctx, sessionID, updateFields); err != nil { + if err = s.chatSessionDAO.UpdateByID(ctx, dao.DB, sessionID, updateFields); err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Session not found!") } return nil, common.CodeServerError, err } - session, err := s.chatSessionDAO.GetByID(ctx, sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) if err != nil { if isChatSessionNotFound(err) { return nil, common.CodeDataError, errors.New("Fail to update a session!") @@ -618,7 +618,7 @@ func (s *ChatSessionService) DeleteSessionMessage(ctx context.Context, userID, c return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - session, err := s.chatSessionDAO.GetByID(ctx, sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) if err != nil || session.DialogID != chatID { if err != nil && !isChatSessionNotFound(err) { return nil, common.CodeServerError, err @@ -663,7 +663,7 @@ func (s *ChatSessionService) DeleteSessionMessage(ctx context.Context, userID, c if err != nil { return nil, common.CodeServerError, err } - if err = s.chatSessionDAO.UpdateByID(ctx, session.ID, map[string]interface{}{ + if err = s.chatSessionDAO.UpdateByID(ctx, dao.DB, session.ID, map[string]interface{}{ "message": messageRaw, "reference": referenceRaw, }); err != nil { @@ -685,7 +685,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, return nil, common.CodeServerError, err } for _, tenantID := range tenantIDs { - exists, err := s.chatSessionDAO.CheckDialogExists(ctx, tenantID, chatID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, dao.DB, tenantID, chatID) if err != nil { return nil, common.CodeServerError, err } @@ -695,7 +695,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, } } if ownerTenantID == "" { - exists, err := s.chatSessionDAO.CheckDialogExists(ctx, userID, chatID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, dao.DB, userID, chatID) if err != nil { return nil, common.CodeServerError, err } @@ -708,7 +708,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, return nil, common.CodeAuthenticationError, errors.New("No authorization.") } - session, err := s.chatSessionDAO.GetByID(ctx, sessionID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) if err != nil || session.DialogID != chatID { if err != nil && !isChatSessionNotFound(err) { return nil, common.CodeServerError, err @@ -779,7 +779,7 @@ func (s *ChatSessionService) UpdateMessageFeedback(ctx context.Context, userID, if err != nil { return nil, common.CodeServerError, err } - if err = s.chatSessionDAO.UpdateByID(ctx, session.ID, map[string]interface{}{"message": messageRaw}); err != nil { + if err = s.chatSessionDAO.UpdateByID(ctx, dao.DB, session.ID, map[string]interface{}{"message": messageRaw}); err != nil { return nil, common.CodeServerError, err } session.Message = messageRaw @@ -1119,7 +1119,7 @@ func (s *ChatSessionService) ensureOwnedChat(ctx context.Context, userID, chatID } for _, tenantID := range tenantIDs { - exists, err := s.chatSessionDAO.CheckDialogExists(ctx, tenantID, chatID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, dao.DB, tenantID, chatID) if err != nil { return false, err } @@ -1128,7 +1128,7 @@ func (s *ChatSessionService) ensureOwnedChat(ctx context.Context, userID, chatID } } - exists, err := s.chatSessionDAO.CheckDialogExists(ctx, userID, chatID) + exists, err := s.chatSessionDAO.CheckDialogExists(ctx, dao.DB, userID, chatID) if err != nil { return false, err } @@ -1274,12 +1274,12 @@ func (s *ChatSessionService) Completion(ctx context.Context, userID string, conv return nil, errors.New("the last content of this conversation is not from user") } - session, err := s.chatSessionDAO.GetByID(ctx, conversationID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, conversationID) if err != nil { return nil, errors.New("Conversation not found") } - dialog, err := s.chatSessionDAO.GetDialogByID(ctx, session.DialogID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, dao.DB, session.DialogID) if err != nil { return nil, errors.New("Dialog not found") } @@ -1361,13 +1361,13 @@ func (s *ChatSessionService) CompletionStream(ctx context.Context, userID string return errors.New("the last content of this conversation is not from user") } - session, err := s.chatSessionDAO.GetByID(ctx, conversationID) + session, err := s.chatSessionDAO.GetByID(ctx, dao.DB, conversationID) if err != nil { streamChan <- fmt.Sprintf("data: %s\n\n", `{"code": 500, "message": "Conversation not found", "data": {"answer": "**ERROR**: Conversation not found", "reference": []}}`) return errors.New("conversation not found") } - dialog, err := s.chatSessionDAO.GetDialogByID(ctx, session.DialogID) + dialog, err := s.chatSessionDAO.GetDialogByID(ctx, dao.DB, session.DialogID) if err != nil { streamChan <- fmt.Sprintf("data: %s\n\n", `{"code": 500, "message": "Dialog not found", "data": {"answer": "**ERROR**: Dialog not found", "reference": []}}`) return errors.New("dialog not found") @@ -1494,12 +1494,12 @@ func (s *ChatSessionService) ChatCompletions( if err = s.checkDialogOwnership(ctx, userID, chatID); err != nil { return fail(err) } - dialog, err = s.chatSessionDAO.GetDialogByID(ctx, chatID) + dialog, err = s.chatSessionDAO.GetDialogByID(ctx, dao.DB, chatID) if err != nil { return fail(errors.New("Chat not found!")) } if sessionID != "" { - session, err = s.chatSessionDAO.GetByID(ctx, sessionID) + session, err = s.chatSessionDAO.GetByID(ctx, dao.DB, sessionID) if err != nil { return fail(errors.New("Session not found!")) } @@ -1829,7 +1829,7 @@ func (s *ChatSessionService) createSessionForCompletion(ctx context.Context, cha UserID: &userID, Reference: refJSON, } - if err := s.chatSessionDAO.Create(ctx, session); err != nil { + if err := s.chatSessionDAO.Create(ctx, dao.DB, session); err != nil { return nil, err } return session, nil @@ -2140,7 +2140,7 @@ func (s *ChatSessionService) updateSessionMessages(ctx context.Context, session "message": messagesJSON, "reference": referenceJSON, } - if err := s.chatSessionDAO.UpdateByID(ctx, session.ID, updates); err != nil { + if err := s.chatSessionDAO.UpdateByID(ctx, dao.DB, session.ID, updates); err != nil { common.Warn("updateSessionMessages: DAO update failed", zap.Error(err)) return } diff --git a/internal/service/chat_session_test.go b/internal/service/chat_session_test.go index 4698b746f6..95f5b9362c 100644 --- a/internal/service/chat_session_test.go +++ b/internal/service/chat_session_test.go @@ -47,7 +47,7 @@ func newFakeSessionStore() *fakeSessionStore { } } -func (f *fakeSessionStore) GetByID(ctx context.Context, id string) (*entity.ChatSession, error) { +func (f *fakeSessionStore) GetByID(ctx context.Context, db *gorm.DB, id string) (*entity.ChatSession, error) { if f.getByIDErr != nil { return nil, f.getByIDErr } @@ -58,8 +58,8 @@ func (f *fakeSessionStore) GetByID(ctx context.Context, id string) (*entity.Chat return s, nil } -func (f *fakeSessionStore) GetBySessionIDAndChatID(ctx context.Context, sessionID, chatID string) (*entity.ChatSession, error) { - s, err := f.GetByID(ctx, sessionID) +func (f *fakeSessionStore) GetBySessionIDAndChatID(ctx context.Context, db *gorm.DB, sessionID, chatID string) (*entity.ChatSession, error) { + s, err := f.GetByID(ctx, db, sessionID) if err != nil { return nil, err } @@ -69,7 +69,7 @@ func (f *fakeSessionStore) GetBySessionIDAndChatID(ctx context.Context, sessionI return s, nil } -func (f *fakeSessionStore) Create(ctx context.Context, conv *entity.ChatSession) error { +func (f *fakeSessionStore) Create(ctx context.Context, db *gorm.DB, conv *entity.ChatSession) error { f.mu.Lock() defer f.mu.Unlock() if f.createErr != nil { @@ -80,7 +80,7 @@ func (f *fakeSessionStore) Create(ctx context.Context, conv *entity.ChatSession) return nil } -func (f *fakeSessionStore) UpdateByID(ctx context.Context, id string, updates map[string]interface{}) error { +func (f *fakeSessionStore) UpdateByID(ctx context.Context, db *gorm.DB, id string, updates map[string]interface{}) error { f.mu.Lock() defer f.mu.Unlock() if f.updateByIDErr != nil { @@ -113,12 +113,12 @@ func (f *fakeSessionStore) UpdateByID(ctx context.Context, id string, updates ma return nil } -func (f *fakeSessionStore) DeleteByID(ctx context.Context, id string) error { +func (f *fakeSessionStore) DeleteByID(ctx context.Context, db *gorm.DB, id string) error { delete(f.sessions, id) return nil } -func (f *fakeSessionStore) ListByChatID(ctx context.Context, chatID string) ([]*entity.ChatSession, error) { +func (f *fakeSessionStore) ListByChatID(ctx context.Context, db *gorm.DB, chatID string) ([]*entity.ChatSession, error) { var result []*entity.ChatSession for _, s := range f.sessions { if s.DialogID == chatID { @@ -128,7 +128,7 @@ func (f *fakeSessionStore) ListByChatID(ctx context.Context, chatID string) ([]* return result, nil } -func (f *fakeSessionStore) GetDialogByID(ctx context.Context, chatID string) (*entity.Chat, error) { +func (f *fakeSessionStore) GetDialogByID(ctx context.Context, db *gorm.DB, chatID string) (*entity.Chat, error) { if f.getDialogErr != nil { return nil, f.getDialogErr } @@ -139,7 +139,7 @@ func (f *fakeSessionStore) GetDialogByID(ctx context.Context, chatID string) (*e return d, nil } -func (f *fakeSessionStore) CheckDialogExists(ctx context.Context, tenantID, chatID string) (bool, error) { +func (f *fakeSessionStore) CheckDialogExists(ctx context.Context, db *gorm.DB, tenantID, chatID string) (bool, error) { key := tenantID + "|" + chatID return f.dialogExists[key], nil } diff --git a/internal/service/connector.go b/internal/service/connector.go index e3af4b6988..45b98fd744 100644 --- a/internal/service/connector.go +++ b/internal/service/connector.go @@ -226,16 +226,16 @@ func (s *ConnectorService) canAccessConnector(connector *entity.Connector, userI } // cancelConnectorTasks Stop connector tasks -func (s *ConnectorService) cancelConnectorTasks(connectorID string) error { - if err := s.connectorDAO.CancelRunningOrScheduledLogs(connectorID); err != nil { +func (s *ConnectorService) cancelConnectorTasks(ctx context.Context, connectorID string) error { + if err := s.connectorDAO.CancelRunningOrScheduledLogs(ctx, dao.DB, connectorID); err != nil { return err } - return s.connectorDAO.UpdateByID(connectorID, map[string]interface{}{"status": string(entity.TaskStatusCancel)}) + return s.connectorDAO.UpdateByID(ctx, dao.DB, connectorID, map[string]interface{}{"status": string(entity.TaskStatusCancel)}) } // CreateConnector creates a connector owned by the current user. // Equivalent to Python's create_connector endpoint. -func (s *ConnectorService) CreateConnector(userID string, req *CreateConnectorRequest) (*entity.Connector, error) { +func (s *ConnectorService) CreateConnector(ctx context.Context, userID string, req *CreateConnectorRequest) (*entity.Connector, error) { refreshFreq := int64(defaultConnectorFreq) if req.RefreshFreq != nil { refreshFreq = *req.RefreshFreq @@ -264,23 +264,23 @@ func (s *ConnectorService) CreateConnector(userID string, req *CreateConnectorRe Status: connectorStatusUnstarted, } - if err := s.connectorDAO.Create(connector); err != nil { + if err := s.connectorDAO.Create(ctx, dao.DB, connector); err != nil { return nil, err } - return s.connectorDAO.GetByID(connector.ID) + return s.connectorDAO.GetByID(ctx, dao.DB, connector.ID) } // GetConnector returns one connector when the user can access its tenant. -func (s *ConnectorService) GetConnector(connectorID, userID string) (*entity.Connector, common.ErrorCode, error) { +func (s *ConnectorService) GetConnector(ctx context.Context, connectorID, userID string) (*entity.Connector, common.ErrorCode, error) { if connectorID == "" { return nil, common.CodeDataError, fmt.Errorf("connector_id is required") } - connector, err := s.connectorDAO.GetByID(connectorID) + connector, err := s.connectorDAO.GetByID(ctx, dao.DB, connectorID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + return nil, common.CodeDataError, fmt.Errorf("can't find this Connector") } return nil, common.CodeServerError, err } @@ -299,18 +299,18 @@ func (s *ConnectorService) GetConnector(connectorID, userID string) (*entity.Con } } - return nil, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + return nil, common.CodeAuthenticationError, fmt.Errorf("no authorization") } // ListConnectors list connectors for a user -func (s *ConnectorService) ListConnectors(userID string) (*ListConnectorsResponse, error) { +func (s *ConnectorService) ListConnectors(ctx context.Context, userID string) (*ListConnectorsResponse, error) { userID = strings.TrimSpace(userID) if userID == "" { return nil, fmt.Errorf("user_id is required") } // Query connectors by tenant ID - connectors, err := s.connectorDAO.ListByTenantID(userID) + connectors, err := s.connectorDAO.ListByTenantID(ctx, dao.DB, userID) if err != nil { return nil, err } @@ -322,8 +322,8 @@ func (s *ConnectorService) ListConnectors(userID string) (*ListConnectorsRespons // accessible reports whether the user can access the connector's tenant. // Mirrors Python's ConnectorService.accessible: owner access plus joined tenants. -func (s *ConnectorService) accessible(connectorID, userID string) (bool, error) { - conn, err := s.connectorDAO.GetByID(connectorID) +func (s *ConnectorService) accessible(ctx context.Context, connectorID, userID string) (bool, error) { + conn, err := s.connectorDAO.GetByID(ctx, dao.DB, connectorID) if err != nil { return false, ErrConnectorNotFound } @@ -351,8 +351,8 @@ func (s *ConnectorService) accessible(connectorID, userID string) (bool, error) // is REST_API (the only source Python currently tests), and that credentials // are present in the stored config. It returns ErrConnectorTestUnsupported for // other sources. -func (s *ConnectorService) TestConnector(connectorID, userID string) error { - ok, err := s.accessible(connectorID, userID) +func (s *ConnectorService) TestConnector(ctx context.Context, connectorID, userID string) error { + ok, err := s.accessible(ctx, connectorID, userID) if err != nil && errors.Is(err, ErrConnectorNotFound) { return ErrConnectorNotFound } @@ -363,7 +363,7 @@ func (s *ConnectorService) TestConnector(connectorID, userID string) error { return ErrConnectorNoAuth } - conn, err := s.connectorDAO.GetByID(connectorID) + conn, err := s.connectorDAO.GetByID(ctx, dao.DB, connectorID) if err != nil { return ErrConnectorNotFound } @@ -383,13 +383,13 @@ func (s *ConnectorService) TestConnector(connectorID, userID string) error { return nil } -func (s *ConnectorService) StartGoogleWebOAuth(userID, source string, req *StartGoogleWebOAuthRequest) (*StartGoogleWebOAuthResponse, common.ErrorCode, error) { +func (s *ConnectorService) StartGoogleWebOAuth(ctx context.Context, userID, source string, req *StartGoogleWebOAuthRequest) (*StartGoogleWebOAuthResponse, common.ErrorCode, error) { source = strings.TrimSpace(source) if source == "" { source = "google-drive" } if source != "google-drive" && source != "gmail" { - return nil, common.CodeArgumentError, fmt.Errorf("Invalid Google OAuth type.") + return nil, common.CodeArgumentError, fmt.Errorf("invalid Google OAuth type") } if req == nil || len(req.Credentials) == 0 { @@ -401,7 +401,7 @@ func (s *ConnectorService) StartGoogleWebOAuth(userID, source string, req *Start redirectURI = defaultGoogleWebOAuthRedirectURI(source) } if redirectURI == "" { - return nil, common.CodeServerError, fmt.Errorf("Google OAuth redirect URI is not configured on the server.") + return nil, common.CodeServerError, fmt.Errorf("not configure Google OAuth redirect URI on the server") } credentials, err := loadGoogleCredentials(req.Credentials) @@ -409,7 +409,7 @@ func (s *ConnectorService) StartGoogleWebOAuth(userID, source string, req *Start return nil, common.CodeArgumentError, err } if hasRefreshToken(credentials) { - return nil, common.CodeArgumentError, fmt.Errorf("Uploaded credentials already include a refresh token.") + return nil, common.CodeArgumentError, fmt.Errorf("uploaded credentials already include a refresh token") } clientConfig, err := getGoogleWebClientConfig(credentials) @@ -424,7 +424,7 @@ func (s *ConnectorService) StartGoogleWebOAuth(userID, source string, req *Start authURI = googleOAuthAuthorizeURL } if clientID == "" || authURI == "" { - return nil, common.CodeServerError, fmt.Errorf("Failed to initialize Google OAuth flow. Please verify the uploaded client configuration.") + return nil, common.CodeServerError, fmt.Errorf("failed to initialize Google OAuth flow. Please verify the uploaded client configuration") } codeVerifier, codeChallenge, err := newPKCEChallenge() @@ -435,12 +435,12 @@ func (s *ConnectorService) StartGoogleWebOAuth(userID, source string, req *Start flowID := utility.GenerateUUID() authorizationURL, err := buildGoogleAuthorizationURL(authURI, clientID, redirectURI, flowID, googleOAuthScopesForSource(source), codeChallenge) if err != nil { - return nil, common.CodeServerError, fmt.Errorf("Failed to initialize Google OAuth flow. Please verify the uploaded client configuration.") + return nil, common.CodeServerError, fmt.Errorf("failed to initialize Google OAuth flow. Please verify the uploaded client configuration") } redisClient := redis.Get() if redisClient == nil { - return nil, common.CodeServerError, fmt.Errorf("Redis is not configured on the server.") + return nil, common.CodeServerError, fmt.Errorf("no configure Redis on the server") } state := googleWebOAuthState{ @@ -451,7 +451,7 @@ func (s *ConnectorService) StartGoogleWebOAuth(userID, source string, req *Start CreatedAt: time.Now().Unix(), } if ok := redisClient.SetObj(webStateCacheKey(flowID, source), state, webFlowTTL); !ok { - return nil, common.CodeServerError, fmt.Errorf("Failed to initialize Google OAuth flow. Please verify the uploaded client configuration.") + return nil, common.CodeServerError, fmt.Errorf("failed to initialize Google OAuth flow. Please verify the uploaded client configuration") } return &StartGoogleWebOAuthResponse{ @@ -461,7 +461,7 @@ func (s *ConnectorService) StartGoogleWebOAuth(userID, source string, req *Start }, common.CodeSuccess, nil } -func (s *ConnectorService) GoogleWebOAuthCallback(source, stateID, oauthError, errorDescription, code string) string { +func (s *ConnectorService) GoogleWebOAuthCallback(ctx context.Context, source, stateID, oauthError, errorDescription, code string) string { source = strings.TrimSpace(source) if source != "google-drive" && source != "gmail" { return renderWebOAuthPopup("", false, "Invalid Google OAuth type.", source) @@ -524,10 +524,10 @@ func (s *ConnectorService) GoogleWebOAuthCallback(source, stateID, oauthError, e return renderWebOAuthPopup(stateID, true, "Authorization completed successfully.", source) } -func (s *ConnectorService) PollGoogleWebOAuthResult(userID, source string, req *PollGoogleWebOAuthResultRequest) (*PollGoogleWebOAuthResultResponse, common.ErrorCode, error) { +func (s *ConnectorService) PollGoogleWebOAuthResult(ctx context.Context, userID, source string, req *PollGoogleWebOAuthResultRequest) (*PollGoogleWebOAuthResultResponse, common.ErrorCode, error) { source = strings.TrimSpace(source) if source != "google-drive" && source != "gmail" { - return nil, common.CodeArgumentError, fmt.Errorf("Invalid Google OAuth type.") + return nil, common.CodeArgumentError, fmt.Errorf("invalid Google OAuth type") } if req == nil || strings.TrimSpace(req.FlowID) == "" { return nil, common.CodeArgumentError, fmt.Errorf("required argument is missing: flow_id") @@ -535,17 +535,17 @@ func (s *ConnectorService) PollGoogleWebOAuthResult(userID, source string, req * redisClient := redis.Get() if redisClient == nil { - return nil, common.CodeRunning, fmt.Errorf("Authorization is still pending.") + return nil, common.CodeRunning, fmt.Errorf("authorization is still pending") } resultKey := webResultCacheKey(strings.TrimSpace(req.FlowID), source) var result googleWebOAuthResult if ok := redisClient.GetObj(resultKey, &result); !ok { - return nil, common.CodeRunning, fmt.Errorf("Authorization is still pending.") + return nil, common.CodeRunning, fmt.Errorf("authorization is still pending") } if result.UserID != userID { - return nil, common.CodePermissionError, fmt.Errorf("You are not allowed to access this authorization result.") + return nil, common.CodePermissionError, fmt.Errorf("you are not allowed to access this authorization result") } redisClient.Delete(resultKey) @@ -582,10 +582,10 @@ func loadGoogleCredentials(raw json.RawMessage) (map[string]interface{}, error) var rawString string if err := json.Unmarshal(raw, &rawString); err != nil { - return nil, fmt.Errorf("Invalid Google credentials JSON.") + return nil, fmt.Errorf("invalid Google credentials JSON") } if err := json.Unmarshal([]byte(strings.TrimSpace(rawString)), &credentials); err != nil || credentials == nil { - return nil, fmt.Errorf("Invalid Google credentials JSON.") + return nil, fmt.Errorf("invalid Google credentials JSON") } return credentials, nil } @@ -604,7 +604,7 @@ func hasRefreshToken(credentials map[string]interface{}) bool { func getGoogleWebClientConfig(credentials map[string]interface{}) (map[string]interface{}, error) { webSection, ok := credentials["web"].(map[string]interface{}) if !ok || webSection == nil { - return nil, fmt.Errorf("Google OAuth JSON must include a 'web' client configuration to use browser-based authorization.") + return nil, fmt.Errorf("google OAuth JSON must include a 'web' client configuration to use browser-based authorization") } return map[string]interface{}{"web": webSection}, nil } @@ -828,28 +828,28 @@ func webOAuthSourceDisplayName(source string) string { return "OAuth" } -func (s *ConnectorService) DeleteConnector(connectorID, userID string) (bool, common.ErrorCode, error) { +func (s *ConnectorService) DeleteConnector(ctx context.Context, connectorID, userID string) (bool, common.ErrorCode, error) { if connectorID == "" { return false, common.CodeDataError, fmt.Errorf("connector_id is required") } - connector, err := s.connectorDAO.GetByID(connectorID) + connector, err := s.connectorDAO.GetByID(ctx, dao.DB, connectorID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return false, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + return false, common.CodeDataError, fmt.Errorf("can't find this Connector") } return false, common.CodeServerError, err } if !s.canAccessConnector(connector, userID) { - return false, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + return false, common.CodeAuthenticationError, fmt.Errorf("no authorization") } - if err = s.cancelConnectorTasks(connector.ID); err != nil { + if err = s.cancelConnectorTasks(ctx, connector.ID); err != nil { return false, common.CodeServerError, err } - if err = s.connectorDAO.DeleteByID(connector.ID); err != nil { + if err = s.connectorDAO.DeleteByID(ctx, dao.DB, connector.ID); err != nil { return false, common.CodeServerError, err } return true, common.CodeSuccess, nil @@ -864,21 +864,21 @@ type UpdateConnectorRequest struct { Status string `json:"status,omitempty"` } -func (s *ConnectorService) UpdateConnector(connectorID, userID string, req *UpdateConnectorRequest) (*entity.Connector, common.ErrorCode, error) { +func (s *ConnectorService) UpdateConnector(ctx context.Context, connectorID, userID string, req *UpdateConnectorRequest) (*entity.Connector, common.ErrorCode, error) { if connectorID == "" { return nil, common.CodeDataError, fmt.Errorf("connector_id is required") } - connector, err := s.connectorDAO.GetByID(connectorID) + connector, err := s.connectorDAO.GetByID(ctx, dao.DB, connectorID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + return nil, common.CodeDataError, fmt.Errorf("can't find this Connector") } return nil, common.CodeServerError, err } if !s.canAccessConnector(connector, userID) { - return nil, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + return nil, common.CodeAuthenticationError, fmt.Errorf("no authorization") } updates := map[string]interface{}{} @@ -898,34 +898,34 @@ func (s *ConnectorService) UpdateConnector(connectorID, userID string, req *Upda } if len(updates) > 0 { - if err := s.connectorDAO.UpdateByID(connectorID, updates); err != nil { + if err = s.connectorDAO.UpdateByID(ctx, dao.DB, connectorID, updates); err != nil { return nil, common.CodeServerError, err } } if req != nil { if req.Reschedule { - if err := s.cancelConnectorTasks(connectorID); err != nil { + if err = s.cancelConnectorTasks(ctx, connectorID); err != nil { return nil, common.CodeServerError, err } - if err := s.connectorDAO.ScheduleConnectorTasks(connectorID); err != nil { + if err = s.connectorDAO.ScheduleConnectorTasks(ctx, dao.DB, connectorID); err != nil { return nil, common.CodeServerError, err } } else if isConnectorCancelStatus(req.Status) { - if err := s.cancelConnectorTasks(connectorID); err != nil { + if err = s.cancelConnectorTasks(ctx, connectorID); err != nil { return nil, common.CodeServerError, err } } else if isConnectorScheduleStatus(req.Status) { - if err := s.connectorDAO.ScheduleConnectorTasks(connectorID); err != nil { + if err = s.connectorDAO.ScheduleConnectorTasks(ctx, dao.DB, connectorID); err != nil { return nil, common.CodeServerError, err } } } - connector, err = s.connectorDAO.GetByID(connectorID) + connector, err = s.connectorDAO.GetByID(ctx, dao.DB, connectorID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + return nil, common.CodeDataError, fmt.Errorf("can't find this Connector") } return nil, common.CodeServerError, err } @@ -944,7 +944,7 @@ func isConnectorScheduleStatus(status string) bool { } // RebuildConnector schedules a rebuild for an accessible connector and knowledge base. -func (s *ConnectorService) RebuildConnector(connectorID, userID, kbID string) (bool, common.ErrorCode, error) { +func (s *ConnectorService) RebuildConnector(ctx context.Context, connectorID, userID, kbID string) (bool, common.ErrorCode, error) { if connectorID == "" { return false, common.CodeDataError, fmt.Errorf("connector_id is required") } @@ -952,33 +952,33 @@ func (s *ConnectorService) RebuildConnector(connectorID, userID, kbID string) (b return false, common.CodeArgumentError, fmt.Errorf("required argument is missing: kb_id") } - connector, err := s.connectorDAO.GetByID(connectorID) + connector, err := s.connectorDAO.GetByID(ctx, dao.DB, connectorID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return false, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + return false, common.CodeDataError, fmt.Errorf("can't find this Connector") } return false, common.CodeServerError, err } if !s.canAccessConnector(connector, userID) { - return false, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + return false, common.CodeAuthenticationError, fmt.Errorf("no authorization") } sourceType := fmt.Sprintf("%s/%s", connector.Source, connector.ID) - documents, err := s.connectorDAO.ListDocumentsByKBAndSourceType(kbID, sourceType) + documents, err := s.connectorDAO.ListDocumentsByKBAndSourceType(ctx, dao.DB, kbID, sourceType) if err != nil { return false, common.CodeServerError, err } - s.deleteConnectorDocumentChunks(connector.TenantID, kbID, documents) + s.deleteConnectorDocumentChunks(ctx, connector.TenantID, kbID, documents) - if err := s.connectorDAO.RebuildConnector(connector, kbID, documents); err != nil { + if err = s.connectorDAO.RebuildConnector(ctx, dao.DB, connector, kbID, documents); err != nil { return false, common.CodeServerError, err } return true, common.CodeSuccess, nil } -func (s *ConnectorService) deleteConnectorDocumentChunks(tenantID, kbID string, documents []*entity.Document) { +func (s *ConnectorService) deleteConnectorDocumentChunks(ctx context.Context, tenantID, kbID string, documents []*entity.Document) { docEngine := engine.Get() if docEngine == nil { return @@ -986,25 +986,25 @@ func (s *ConnectorService) deleteConnectorDocumentChunks(tenantID, kbID string, indexName := fmt.Sprintf("ragflow_%s", tenantID) for _, document := range documents { - _, _ = docEngine.DeleteChunks(context.Background(), map[string]interface{}{"doc_id": document.ID}, indexName, kbID) + _, _ = docEngine.DeleteChunks(ctx, map[string]interface{}{"doc_id": document.ID}, indexName, kbID) } } -func (s *ConnectorService) ListLog(connectorID, userID string, page, pageSize int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) { +func (s *ConnectorService) ListLog(ctx context.Context, connectorID, userID string, page, pageSize int) ([]*entity.ConnectorSyncLog, int64, common.ErrorCode, error) { if connectorID == "" { return nil, 0, common.CodeDataError, fmt.Errorf("connector_id is required") } - connector, err := s.connectorDAO.GetByID(connectorID) + connector, err := s.connectorDAO.GetByID(ctx, dao.DB, connectorID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - return nil, 0, common.CodeDataError, fmt.Errorf("Can't find this Connector!") + return nil, 0, common.CodeDataError, fmt.Errorf("can't find this Connector") } return nil, 0, common.CodeServerError, err } if !s.canAccessConnector(connector, userID) { - return nil, 0, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + return nil, 0, common.CodeAuthenticationError, fmt.Errorf("no authorization") } if page < 1 { @@ -1015,7 +1015,7 @@ func (s *ConnectorService) ListLog(connectorID, userID string, page, pageSize in } offset := (page - 1) * pageSize - logs, total, err := s.connectorDAO.ListLogsByConnectorID(connectorID, offset, pageSize) + logs, total, err := s.connectorDAO.ListLogsByConnectorID(ctx, dao.DB, connectorID, offset, pageSize) if err != nil { return nil, 0, common.CodeServerError, fmt.Errorf("failed to fetch connector logs: %w", err) } @@ -1025,7 +1025,7 @@ func (s *ConnectorService) ListLog(connectorID, userID string, page, pageSize in return logs, total, common.CodeSuccess, nil } -func (s *ConnectorService) StartBoxWebOAuth(userID string, req *StartBoxWebOAuthRequest) (*StartBoxWebOAuthResponse, common.ErrorCode, error) { +func (s *ConnectorService) StartBoxWebOAuth(ctx context.Context, userID string, req *StartBoxWebOAuthRequest) (*StartBoxWebOAuthResponse, common.ErrorCode, error) { var clientID, clientSecret, redirectURI string if req != nil { clientID = strings.TrimSpace(req.ClientID) @@ -1033,7 +1033,7 @@ func (s *ConnectorService) StartBoxWebOAuth(userID string, req *StartBoxWebOAuth redirectURI = strings.TrimSpace(req.RedirectURI) } if clientID == "" || clientSecret == "" { - return nil, common.CodeArgumentError, fmt.Errorf("Box client_id and client_secret are required.") + return nil, common.CodeArgumentError, fmt.Errorf("box client_id and client_secret are required") } if redirectURI == "" { redirectURI = defaultBoxWebOAuthRedirectURI() @@ -1047,7 +1047,7 @@ func (s *ConnectorService) StartBoxWebOAuth(userID string, req *StartBoxWebOAuth redisClient := connectorRedisGet() if redisClient == nil { - return nil, common.CodeServerError, fmt.Errorf("Redis is not configured on the server.") + return nil, common.CodeServerError, fmt.Errorf("not connected Redis client on the server") } state := boxWebOAuthState{ @@ -1059,7 +1059,7 @@ func (s *ConnectorService) StartBoxWebOAuth(userID string, req *StartBoxWebOAuth CreatedAt: time.Now().Unix(), } if ok := redisClient.SetObj(webStateCacheKey(flowID, "box"), state, webFlowTTL); !ok { - return nil, common.CodeServerError, fmt.Errorf("Failed to initialize Box OAuth flow. Please verify the client configuration.") + return nil, common.CodeServerError, fmt.Errorf("failed to initialize Box OAuth flow. Please verify the client configuration") } return &StartBoxWebOAuthResponse{ @@ -1069,7 +1069,7 @@ func (s *ConnectorService) StartBoxWebOAuth(userID string, req *StartBoxWebOAuth }, common.CodeSuccess, nil } -func (s *ConnectorService) BoxWebOAuthCallback(flowID string, oauthError string, errorDescription string, code string) string { +func (s *ConnectorService) BoxWebOAuthCallback(ctx context.Context, flowID string, oauthError string, errorDescription string, code string) string { flowID = strings.TrimSpace(flowID) if flowID == "" { return renderWebOAuthPopup("", false, "Missing OAuth parameters.", "box") @@ -1125,24 +1125,24 @@ func (s *ConnectorService) BoxWebOAuthCallback(flowID string, oauthError string, return renderWebOAuthPopup(flowID, true, "Authorization completed successfully.", "box") } -func (s *ConnectorService) PollBoxWebOAuthResult(userID string, req *PollBoxWebOAuthResultRequest) (*PollBoxWebOAuthResultResponse, common.ErrorCode, error) { +func (s *ConnectorService) PollBoxWebOAuthResult(ctx context.Context, userID string, req *PollBoxWebOAuthResultRequest) (*PollBoxWebOAuthResultResponse, common.ErrorCode, error) { if req == nil || strings.TrimSpace(req.FlowID) == "" { return nil, common.CodeArgumentError, fmt.Errorf("required argument is missing: flow_id") } redisClient := connectorRedisGet() if redisClient == nil { - return nil, common.CodeRunning, fmt.Errorf("Authorization is still pending.") + return nil, common.CodeRunning, fmt.Errorf("authorization is still pending") } resultKey := webResultCacheKey(strings.TrimSpace(req.FlowID), "box") var result boxWebOAuthCredentials if ok := redisClient.GetObj(resultKey, &result); !ok { - return nil, common.CodeRunning, fmt.Errorf("Authorization is still pending.") + return nil, common.CodeRunning, fmt.Errorf("authorization is still pending") } if result.UserID != userID { - return nil, common.CodePermissionError, fmt.Errorf("You are not allowed to access this authorization result.") + return nil, common.CodePermissionError, fmt.Errorf("you are not allowed to access this authorization result") } redisClient.Delete(resultKey) @@ -1203,7 +1203,7 @@ func exchangeBoxAuthorizationCode(clientID string, clientSecret string, redirect } var token boxOAuthTokenResponse - if err := json.Unmarshal(body, &token); err != nil { + if err = json.Unmarshal(body, &token); err != nil { return nil, err } if resp.StatusCode >= http.StatusBadRequest || token.Error != "" { diff --git a/internal/service/connector_box_oauth_test.go b/internal/service/connector_box_oauth_test.go index b31af12b31..d1d8a1c61a 100644 --- a/internal/service/connector_box_oauth_test.go +++ b/internal/service/connector_box_oauth_test.go @@ -140,15 +140,16 @@ func TestExchangeBoxAuthorizationCodeRequestBody(t *testing.T) { func TestStartBoxWebOAuthRequiresClientCredentials(t *testing.T) { svc := NewConnectorService() + ctx := t.Context() - _, code, err := svc.StartBoxWebOAuth("user-1", &StartBoxWebOAuthRequest{ClientID: "client-1"}) + _, code, err := svc.StartBoxWebOAuth(ctx, "user-1", &StartBoxWebOAuthRequest{ClientID: "client-1"}) if err == nil { t.Fatal("expected error") } if code != common.CodeArgumentError { t.Fatalf("code=%v want=%v", code, common.CodeArgumentError) } - if err.Error() != "Box client_id and client_secret are required." { + if err.Error() != "box client_id and client_secret are required" { t.Fatalf("error=%q", err.Error()) } } @@ -157,15 +158,16 @@ func TestPollBoxWebOAuthResultPendingWithoutRedis(t *testing.T) { forceNilConnectorRedis(t) svc := NewConnectorService() + ctx := t.Context() - _, code, err := svc.PollBoxWebOAuthResult("user-1", &PollBoxWebOAuthResultRequest{FlowID: "flow-1"}) + _, code, err := svc.PollBoxWebOAuthResult(ctx, "user-1", &PollBoxWebOAuthResultRequest{FlowID: "flow-1"}) if err == nil { t.Fatal("expected pending error") } if code != common.CodeRunning { t.Fatalf("code=%v want=%v", code, common.CodeRunning) } - if err.Error() != "Authorization is still pending." { + if err.Error() != "authorization is still pending" { t.Fatalf("error=%q", err.Error()) } } diff --git a/internal/service/dataset/crud.go b/internal/service/dataset/crud.go index 6bfe399b64..5fd07e5b07 100644 --- a/internal/service/dataset/crud.go +++ b/internal/service/dataset/crud.go @@ -155,10 +155,10 @@ func (d *DatasetService) CreateDataset(req *service.CreateDatasetRequest, tenant return datasetToMap(createdKB), common.CodeSuccess, nil } -func (d *DatasetService) GetDataset(datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) GetDataset(ctx context.Context, datasetID, userID string) (map[string]interface{}, common.ErrorCode, error) { datasetID = strings.TrimSpace(datasetID) if datasetID == "" { - return nil, common.CodeDataError, errors.New("Lack of \"Dataset ID\"") + return nil, common.CodeDataError, errors.New("lack of \"Dataset ID\"") } normalizedID, err := normalizeDatasetID(datasetID) @@ -168,25 +168,25 @@ func (d *DatasetService) GetDataset(datasetID, userID string) (map[string]interf datasetID = normalizedID if !d.kbDAO.Accessible(datasetID, userID) { - return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID) + return nil, common.CodeDataError, fmt.Errorf("user '%s' lacks permission for dataset '%s'", userID, datasetID) } kb, err := d.kbDAO.GetByID(datasetID) if err != nil || kb == nil { - return nil, common.CodeDataError, errors.New("Invalid Dataset ID") + return nil, common.CodeDataError, errors.New("invalid Dataset ID") } data := datasetToMap(kb) size, err := d.documentDAO.SumSizeByDatasetID(datasetID) if err != nil { - return nil, common.CodeServerError, errors.New("Database operation failed") + return nil, common.CodeServerError, errors.New("database operation failed") } data["size"] = size - connectors, err := d.connectorDAO.ListByDatasetID(datasetID) + connectors, err := d.connectorDAO.ListByDatasetID(ctx, dao.DB, datasetID) if err != nil { - return nil, common.CodeServerError, errors.New("Database operation failed") + return nil, common.CodeServerError, errors.New("database operation failed") } data["connectors"] = datasetConnectorsOrEmpty(connectors) diff --git a/internal/service/dataset/update.go b/internal/service/dataset/update.go index e869c0e49f..cde3a7113a 100644 --- a/internal/service/dataset/update.go +++ b/internal/service/dataset/update.go @@ -26,14 +26,14 @@ type datasetPagerankUpdate struct { datasetID string } -func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) { +func (d *DatasetService) UpdateDataset(ctx context.Context, datasetID, tenantID string, req service.UpdateDatasetRequest) (map[string]interface{}, common.ErrorCode, error) { datasetID = strings.TrimSpace(datasetID) tenantID = strings.TrimSpace(tenantID) if _, err := d.kbDAO.GetByID(datasetID); err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("Dataset not found") + return nil, common.CodeDataError, errors.New("dataset not found") } - return nil, common.CodeServerError, errors.New("Database operation failed") + return nil, common.CodeServerError, errors.New("database operation failed") } connectorsProvided := req.Connectors != nil @@ -61,7 +61,7 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U if req.Name != nil { name := strings.TrimSpace(*req.Name) if name == "" { - return nil, common.CodeDataError, errors.New("`name` is required.") + return nil, common.CodeDataError, errors.New("`name` is required") } if len(name) > 128 { return nil, common.CodeDataError, errors.New("String should have at most 128 characters") @@ -98,6 +98,10 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U simpleUpdates["permission"] = permission } + if req.ParseType == nil && req.ParserID != nil && req.PipelineID != nil { + return nil, common.CodeDataError, errors.New("mutually exclusive") + } + if req.ParserID != nil || req.PipelineID != nil || req.ParseType != nil { isBuiltin, isPipeline, modeErr := service.ValidateParseTypeMode(req.ParseType, req.ParserID, req.PipelineID) if modeErr != nil { @@ -147,7 +151,7 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U return nil, common.CodeDataError, errors.New("Input should be less than or equal to 100") } if d.docEngine == nil { - return nil, common.CodeServerError, errors.New("Document engine is not initialized") + return nil, common.CodeServerError, errors.New("document engine is not initialized") } if !d.docEngine.SupportsPageRank() { return nil, common.CodeDataError, errors.New("'pagerank' can only be set when doc_engine is elasticsearch") @@ -157,7 +161,7 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U requestedAnyUpdate := len(simpleUpdates) > 0 || connectorsProvided || parserIDProvided || pipelineID != nil || embdIDProvided || pagerankRequested || req.ParserConfig != nil || req.ParserConfigProvided if !requestedAnyUpdate { - return nil, common.CodeDataError, errors.New("No properties were modified") + return nil, common.CodeDataError, errors.New("no properties were modified") } txCode := common.CodeSuccess @@ -173,7 +177,7 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U if req.Permission != nil && lockedKB.TenantID != tenantID { txCode = common.CodeDataError - return errors.New("Only dataset owner can change permission") + return errors.New("only dataset owner can change permission") } updates := make(map[string]interface{}, len(simpleUpdates)+6) @@ -186,11 +190,11 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U lookupErr := tx.Where("LOWER(name) = LOWER(?) AND tenant_id = ? AND status = ?", nameValue, tenantID, string(entity.StatusValid)).First(&existing).Error if lookupErr != nil && !dao.IsNotFoundErr(lookupErr) { txCode = common.CodeServerError - return errors.New("Database operation failed") + return errors.New("database operation failed") } if lookupErr == nil { txCode = common.CodeDataError - return fmt.Errorf("Dataset name '%s' already exists", nameValue) + return fmt.Errorf("dataset name '%s' already exists", nameValue) } } @@ -284,37 +288,37 @@ func (d *DatasetService) UpdateDataset(datasetID, tenantID string, req service.U if dao.IsDuplicateKeyErr(err) { if nameValue, ok := updates["name"].(string); ok { txCode = common.CodeDataError - return fmt.Errorf("Dataset name '%s' already exists", nameValue) + return fmt.Errorf("dataset name '%s' already exists", nameValue) } txCode = common.CodeDataError - return errors.New("Dataset name already exists") + return errors.New("dataset name already exists") } txCode = common.CodeServerError - return errors.New("Update dataset error.(Database error)") + return errors.New("dataset update error. (database error)") } } if connectorsProvided { - if err = d.connectorDAO.LinkDatasetConnectorsTx(tx, lockedKB.ID, lockedKB.TenantID, connectorLinks); err != nil { + if err = d.connectorDAO.LinkDatasetConnectorsTx(ctx, tx, lockedKB.ID, lockedKB.TenantID, connectorLinks); err != nil { if dao.IsConnectorNotAccessibleErr(err) { txCode = common.CodeDataError return err } txCode = common.CodeServerError - return errors.New("Database operation failed") + return errors.New("database operation failed") } } updatedKB = &entity.Knowledgebase{} if err = tx.Where("id = ? AND status = ?", lockedKB.ID, string(entity.StatusValid)).First(updatedKB).Error; err != nil { txCode = common.CodeDataError - return errors.New("Dataset updated failed") + return errors.New("dataset updated failed") } - linkedConnectors, err = d.connectorDAO.ListByDatasetIDTx(tx, lockedKB.ID) + linkedConnectors, err = d.connectorDAO.ListByDatasetIDTx(ctx, tx, lockedKB.ID) if err != nil { txCode = common.CodeServerError - return errors.New("Database operation failed") + return errors.New("database operation failed") } return nil @@ -353,16 +357,16 @@ func (d *DatasetService) lockAccessibleDatasetForUpdate(tx *gorm.DB, datasetID, First(&kb).Error if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, errors.New("Dataset not found") + return nil, common.CodeDataError, errors.New("dataset not found") } - return nil, common.CodeServerError, errors.New("Database operation failed") + return nil, common.CodeServerError, errors.New("database operation failed") } if kb.TenantID == userID { return &kb, common.CodeSuccess, nil } if kb.Permission != string(entity.TenantPermissionTeam) { - return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID) + return nil, common.CodeDataError, fmt.Errorf("user '%s' lacks permission for dataset '%s'", userID, datasetID) } var relation entity.UserTenant @@ -371,9 +375,9 @@ func (d *DatasetService) lockAccessibleDatasetForUpdate(tx *gorm.DB, datasetID, First(&relation).Error if err != nil { if dao.IsNotFoundErr(err) { - return nil, common.CodeDataError, fmt.Errorf("User '%s' lacks permission for dataset '%s'", userID, datasetID) + return nil, common.CodeDataError, fmt.Errorf("user '%s' lacks permission for dataset '%s'", userID, datasetID) } - return nil, common.CodeServerError, errors.New("Database operation failed") + return nil, common.CodeServerError, errors.New("database operation failed") } return &kb, common.CodeSuccess, nil diff --git a/internal/service/dataset/update_test.go b/internal/service/dataset/update_test.go index 3171561efa..fe1a0a1b08 100644 --- a/internal/service/dataset/update_test.go +++ b/internal/service/dataset/update_test.go @@ -43,7 +43,8 @@ func TestDatasetServiceUpdateDatasetUpdatesFields(t *testing.T) { embeddingModel := "BAAI/bge-large-zh-v1.5@Builtin" parseType := 1 - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ Name: &name, Description: &description, Language: &language, @@ -96,6 +97,30 @@ func TestDatasetServiceUpdateDatasetUpdatesFields(t *testing.T) { } } +func TestUpdateDataset_RejectsSimultaneousParserIDAndPipelineID(t *testing.T) { + db := setupDatasetUpdateTestDB(t) + pushServiceDB(t, db) + insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + + chunkMethod := "book" + pipelineID := "abcdef0123456789abcdef0123456789" + + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ + ParserID: &chunkMethod, + PipelineID: &pipelineID, + }) + if err == nil { + t.Fatal("expected mutual-exclusivity error when both parser_id and pipeline_id are set") + } + if code != common.CodeDataError { + t.Fatalf("expected data error code, got %d", code) + } + if !strings.Contains(err.Error(), "mutually exclusive") { + t.Fatalf("expected error to mention 'mutually exclusive', got: %v", err) + } +} + func TestUpdateDataset_ParseTypeBuiltinClearsPipelineID(t *testing.T) { db := setupDatasetUpdateTestDB(t) pushServiceDB(t, db) @@ -107,7 +132,8 @@ func TestUpdateDataset_ParseTypeBuiltinClearsPipelineID(t *testing.T) { pipelineID := "ABCDEF0123456789ABCDEF0123456789" parseType := 1 - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ ParserID: &chunkMethod, PipelineID: &pipelineID, ParseType: &parseType, @@ -138,7 +164,8 @@ func TestUpdateDataset_ParseTypePipelineIgnoresParserID(t *testing.T) { pipelineID := "ABCDEF0123456789ABCDEF0123456789" parseType := 2 - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ ParserID: &chunkMethod, PipelineID: &pipelineID, ParseType: &parseType, @@ -178,7 +205,8 @@ func TestUpdateDataset_ParseTypePipelineCleansConfigAgainstCanvas(t *testing.T) "Parser:CustomP": map[string]interface{}{"chunk_token_num": float64(256)}, } - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ ParserID: &chunkMethod, PipelineID: &pipelineID, ParseType: &parseType, @@ -221,7 +249,8 @@ func TestUpdateDatasetRejectsInvalidPages(t *testing.T) { chunkMethod := "manual" parseType := 1 - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ ParserID: &chunkMethod, ParseType: &parseType, ParserConfig: map[string]interface{}{ @@ -246,7 +275,8 @@ func TestDatasetServiceGetDatasetReturnsEmptyConnectorList(t *testing.T) { datasetID := "11111111111141118111111111111111" insertDatasetUpdateKB(t, datasetID, "tenant-1", "Original") - result, code, err := testDatasetUpdateService(t).GetDataset("11111111-1111-4111-8111-111111111111", "tenant-1") + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).GetDataset(ctx, "11111111-1111-4111-8111-111111111111", "tenant-1") if err != nil { t.Fatalf("GetDataset failed: %v", err) } @@ -270,14 +300,15 @@ func TestDatasetServiceUpdateDatasetRejectsMissingDataset(t *testing.T) { pushServiceDB(t, db) name := "Renamed" - _, code, err := testDatasetUpdateService(t).UpdateDataset("missing-kb", "tenant-1", service.UpdateDatasetRequest{Name: &name}) + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "missing-kb", "tenant-1", service.UpdateDatasetRequest{Name: &name}) if err == nil { t.Fatal("expected missing dataset error") } if code != common.CodeDataError { t.Fatalf("expected data error code, got %d", code) } - if err.Error() != "Dataset not found" { + if err.Error() != "dataset not found" { t.Fatalf("unexpected error: %v", err) } } @@ -288,7 +319,8 @@ func TestDatasetServiceUpdateDatasetRejectsNonOwner(t *testing.T) { insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") name := "Renamed" - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-2", service.UpdateDatasetRequest{Name: &name}) + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-2", service.UpdateDatasetRequest{Name: &name}) if err == nil { t.Fatal("expected permission error") } @@ -312,7 +344,8 @@ func TestDatasetServiceUpdateDatasetRejectsTeamMemberPermissionChange(t *testing insertDatasetUpdateTeamMember(t, "user-1", "owner-1") permission := string(entity.TenantPermissionMe) - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "user-1", service.UpdateDatasetRequest{ + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "user-1", service.UpdateDatasetRequest{ Permission: &permission, }) if err == nil { @@ -321,7 +354,7 @@ func TestDatasetServiceUpdateDatasetRejectsTeamMemberPermissionChange(t *testing if code != common.CodeDataError { t.Fatalf("expected data error code, got %d", code) } - if err.Error() != "Only dataset owner can change permission" { + if err.Error() != "only dataset owner can change permission" { t.Fatalf("unexpected error: %v", err) } @@ -340,14 +373,15 @@ func TestDatasetServiceUpdateDatasetValidatesName(t *testing.T) { insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") name := " " - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{Name: &name}) + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{Name: &name}) if err == nil { t.Fatal("expected name validation error") } if code != common.CodeDataError { t.Fatalf("expected data error code, got %d", code) } - if err.Error() != "`name` is required." { + if err.Error() != "`name` is required" { t.Fatalf("unexpected error: %v", err) } } @@ -359,7 +393,8 @@ func TestDatasetServiceUpdateDatasetRejectsDuplicateName(t *testing.T) { insertDatasetUpdateKB(t, "kb-2", "tenant-1", "Existing") name := "Existing" - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{Name: &name}) + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{Name: &name}) if err == nil { t.Fatal("expected duplicate name error") } @@ -376,14 +411,15 @@ func TestDatasetServiceUpdateDatasetRejectsNoPropertiesModified(t *testing.T) { pushServiceDB(t, db) insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{}) + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{}) if err == nil { t.Fatal("expected no-op update error") } if code != common.CodeDataError { t.Fatalf("expected data error code, got %d", code) } - if err.Error() != "No properties were modified" { + if err.Error() != "no properties were modified" { t.Fatalf("unexpected error: %v", err) } } @@ -395,7 +431,8 @@ func TestDatasetServiceUpdateDatasetLinksConnectors(t *testing.T) { insertDatasetUpdateConnector(t, "connector-1", "tenant-1") autoParse := "0" - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ Connectors: &[]service.DatasetConnectorRequest{{ID: "connector-1", AutoParse: autoParse}}, }) if err != nil { @@ -431,7 +468,8 @@ func TestDatasetServiceUpdateDatasetRejectsCrossTenantConnector(t *testing.T) { insertDatasetUpdateConnector(t, "connector-1", "tenant-2") autoParse := "0" - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ Connectors: &[]service.DatasetConnectorRequest{{ID: "connector-1", AutoParse: autoParse}}, }) if err == nil { @@ -464,7 +502,8 @@ func TestDatasetServiceUpdateDatasetAcceptsProviderInstanceEmbedding(t *testing. insertDatasetUpdateTenantModel(t, "model-1", "provider-1", "instance-1", "embedding-2", int(entity.ModelTypeEmbedding)) embeddingModel := "embedding-2@test@ZHIPU-AI" - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ EmbeddingModel: &embeddingModel, }) if err != nil { @@ -494,8 +533,9 @@ func TestDatasetServiceUpdateDatasetAcceptsEmbeddingModelID(t *testing.T) { insertDatasetUpdateModelInstance(t, "instance-1", "provider-1", "test") insertDatasetUpdateTenantModel(t, "aabbccdd11223344aabbccdd11223344", "provider-1", "instance-1", "embedding-2", int(entity.ModelTypeEmbedding)) + ctx := t.Context() embeddingModelID := "aabbccdd11223344aabbccdd11223344" - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ EmbeddingModel: &embeddingModelID, }) if err != nil { @@ -525,8 +565,9 @@ func TestDatasetServiceUpdateDatasetRejectsEmptyConnectorID(t *testing.T) { pushServiceDB(t, db) insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") + ctx := t.Context() connectors := []service.DatasetConnectorRequest{{ID: " "}} - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ Connectors: &connectors, }) if err == nil { @@ -557,11 +598,12 @@ func TestDatasetServiceUpdateDatasetRejectsInvalidEmbeddingModelFormat(t *testin {"empty_provider", "BAAI/bge-small-en-v1.5@", "Both model_name and provider must be non-empty strings"}, } + ctx := t.Context() svc := testDatasetUpdateService(t) for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { embdModel := tc.embeddingModel - _, code, err := svc.UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + _, code, err := svc.UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ EmbeddingModel: &embdModel, }) if err == nil { @@ -583,8 +625,9 @@ func TestDatasetServiceUpdateDatasetRejectsDuplicateNameCaseInsensitive(t *testi insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") insertDatasetUpdateKB(t, "kb-2", "tenant-1", "Existing") + ctx := t.Context() uppercaseName := "EXISTING" - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ Name: &uppercaseName, }) if err == nil { @@ -610,8 +653,9 @@ func TestDatasetServiceUpdateDatasetPreservesUnmodifiedFields(t *testing.T) { "language": language, }) + ctx := t.Context() newName := "Renamed Only" - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ Name: &newName, }) if err != nil { @@ -654,8 +698,9 @@ func TestDatasetServiceUpdateDatasetPreservesParserConfigOnEmptyUpdate(t *testin "delimiter": "\n", }) + ctx := t.Context() name := "Updated Name" - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ Name: &name, }) if err != nil { @@ -925,7 +970,8 @@ func TestUpdateDataset_StripsUnknownParam_Builtin(t *testing.T) { pushServiceDB(t, db) insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ ParserConfig: map[string]interface{}{ "Parser:HipSignsRhyme": map[string]interface{}{ "no_such_param": 1, @@ -961,7 +1007,8 @@ func TestUpdateDataset_AcceptsValidComponentParams_Builtin(t *testing.T) { pushServiceDB(t, db) insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original") - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, 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"}, @@ -1003,7 +1050,8 @@ func TestUpdateDataset_StripsCanvasUnknownParam(t *testing.T) { seedDatasetUpdateCanvas(t, "canvas-1", "tenant-1", dsl) insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original", "canvas-1") - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ ParserConfig: map[string]interface{}{ "Parser:NoSuch": map[string]interface{}{ "pdf": map[string]interface{}{}, @@ -1041,7 +1089,8 @@ func TestUpdateDataset_AcceptsValidCanvasComponentParams(t *testing.T) { seedDatasetUpdateCanvas(t, "canvas-1", "tenant-1", dsl) insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original", "canvas-1") - result, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + ctx := t.Context() + result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ ParserConfig: map[string]interface{}{ "Parser:CustomRhyme": map[string]interface{}{ "pdf": map[string]interface{}{}, @@ -1071,8 +1120,9 @@ func TestUpdateDataset_SwitchCanvasToBuiltinValidatesAgainstBuiltin(t *testing.T insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original", "canvas-1") chunkMethod := "naive" + ctx := t.Context() parseType := 1 - _, code, err := testDatasetUpdateService(t).UpdateDataset("kb-1", "tenant-1", service.UpdateDatasetRequest{ + _, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", "tenant-1", service.UpdateDatasetRequest{ ParserID: &chunkMethod, ParseType: &parseType, ParserConfig: map[string]interface{}{ diff --git a/internal/service/llm.go b/internal/service/llm.go index 954f7b8e71..0a51a8b09a 100644 --- a/internal/service/llm.go +++ b/internal/service/llm.go @@ -25,8 +25,6 @@ import ( "ragflow/internal/dao" ) -var DB = dao.DB - // LLMService LLM service type LLMService struct { tenantLLMDAO *dao.TenantLLMDAO @@ -382,7 +380,7 @@ func (s *LLMService) SetAPIKey(tenantID string, req *SetAPIKeyRequest) (*SetAPIK "api_base": baseURL, "max_tokens": maxTokens, } - DB.Model(&entity.TenantLLM{}). + dao.DB.Model(&entity.TenantLLM{}). Where("tenant_id = ? AND llm_factory = ? AND llm_name = ?", tenantID, factory, llm.LLMName). Updates(updates) } else { diff --git a/internal/service/stats.go b/internal/service/stats.go index e9e08eb347..14750f5772 100644 --- a/internal/service/stats.go +++ b/internal/service/stats.go @@ -57,7 +57,7 @@ func (s *StatsService) GetStats(ctx context.Context, userID, fromDate, toDate st return nil, ErrTenantNotFound } - rows, err := dao.NewAPI4ConversationDAO().Stats(ctx, tenants[0].TenantID, fromDate, toDate, source) + rows, err := dao.NewAPI4ConversationDAO().Stats(ctx, dao.DB, tenants[0].TenantID, fromDate, toDate, source) if err != nil { return nil, err } diff --git a/internal/service/toc_enhancer.go b/internal/service/toc_enhancer.go index 3f41496f3b..dfc5767b02 100644 --- a/internal/service/toc_enhancer.go +++ b/internal/service/toc_enhancer.go @@ -279,7 +279,7 @@ func (e *TOCEnhancer) Enhance(ctx context.Context, kbinfos map[string]interface{ } indexNames := make([]string, 0, len(e.tenantIDs)) for _, tid := range e.tenantIDs { - indexNames = append(indexNames, indexName(tid)) + indexNames = append(indexNames, getIndexName(tid)) } tocResp, err := e.docEngine.Search(ctx, &types.SearchRequest{ IndexNames: indexNames, @@ -547,7 +547,7 @@ func (e *TOCEnhancer) fetchChunk(ctx context.Context, chunkID, docID, kbID strin } indexNames := make([]string, 0, len(e.tenantIDs)) for _, tid := range e.tenantIDs { - indexNames = append(indexNames, indexName(tid)) + indexNames = append(indexNames, getIndexName(tid)) } resp, err := e.docEngine.Search(ctx, &types.SearchRequest{ IndexNames: indexNames, @@ -563,8 +563,8 @@ func (e *TOCEnhancer) fetchChunk(ctx context.Context, chunkID, docID, kbID strin return resp.Chunks[0], nil } -// indexName returns the search index name for a tenant. -func indexName(tenantID string) string { +// getIndexName returns the search index name for a tenant. +func getIndexName(tenantID string) string { return "ragflow_" + tenantID } diff --git a/internal/service/toc_enhancer_test.go b/internal/service/toc_enhancer_test.go index ff69123713..faf65b77fe 100644 --- a/internal/service/toc_enhancer_test.go +++ b/internal/service/toc_enhancer_test.go @@ -308,7 +308,7 @@ func TestSortAndTrimChunks_AllKept(t *testing.T) { } func TestIndexName(t *testing.T) { - if got := indexName("tenant1"); got != "ragflow_tenant1" { + if got := getIndexName("tenant1"); got != "ragflow_tenant1" { t.Errorf("expected 'ragflow_tenant1', got %q", got) } } diff --git a/internal/service/user.go b/internal/service/user.go index d49ffa5e33..fbef5d8e7c 100644 --- a/internal/service/user.go +++ b/internal/service/user.go @@ -1020,7 +1020,7 @@ func (s *UserService) GetUserByAPIToken(ctx context.Context, authorization strin // Query API token from database apiTokenDAO := dao.NewAPITokenDAO() - userToken, err := apiTokenDAO.GetUserByAPIToken(ctx, token) + userToken, err := apiTokenDAO.GetUserByAPIToken(ctx, dao.DB, token) if err != nil { return nil, common.CodeUnauthorized, fmt.Errorf("invalid access token") } @@ -1066,7 +1066,7 @@ func (s *UserService) GetAPITokenByBeta(ctx context.Context, authorization strin return nil, fmt.Errorf("invalid authorization format") } apiTokenDAO := dao.NewAPITokenDAO() - tokens, err := apiTokenDAO.GetByBeta(ctx, token) + tokens, err := apiTokenDAO.GetByBeta(ctx, dao.DB, token) if err != nil { return nil, err } @@ -1101,7 +1101,7 @@ func (s *UserService) GetUserByBetaAPIToken(ctx context.Context, authorization s } apiTokenDAO := dao.NewAPITokenDAO() - userTokens, err := apiTokenDAO.GetByBeta(ctx, token) + userTokens, err := apiTokenDAO.GetByBeta(ctx, dao.DB, token) if err != nil || len(userTokens) == 0 { return nil, common.CodeUnauthorized, fmt.Errorf("invalid beta access token") } diff --git a/test/testcases/restful_api/test_chats.py b/test/testcases/restful_api/test_chats.py index d860b85748..16b8c938e6 100644 --- a/test/testcases/restful_api/test_chats.py +++ b/test/testcases/restful_api/test_chats.py @@ -131,8 +131,8 @@ def test_chat_crud_cycle(rest_client, clear_chats): @pytest.mark.parametrize( "name, expected_fragment", [ - ("", "`name` is required."), - (" ", "`name` is required."), + ("", "`name` is required"), + (" ", "`name` is required"), ], ) def test_chat_create_name_validation(rest_client, clear_chats, name, expected_fragment): @@ -154,7 +154,7 @@ def test_chat_duplicate_name_validation(rest_client, clear_chats): assert second.status_code == 200 second_payload = second.json() assert second_payload["code"] == 102, second_payload - assert "Duplicated chat name" in second_payload["message"], second_payload + assert "duplicated chat name" in second_payload["message"], second_payload @pytest.mark.p2 @@ -1635,7 +1635,7 @@ def test_chat_create_prompt_contract(rest_client, clear_chats): @pytest.mark.p2 def test_chat_create_additional_guards_contract(rest_client, clear_chats): cases = [ - ("reject tenant_id override", {"tenant_id": "tenant-should-not-pass"}, "`tenant_id` must not be provided."), + ("reject tenant_id override", {"tenant_id": "tenant-should-not-pass"}, "`tenant_id` must not be provided"), ("reject unknown rerank_id", {"rerank_id": "unknown-rerank-model"}, "`rerank_id` unknown-rerank-model doesn't exist"), ] @@ -1682,13 +1682,13 @@ def test_chat_update_name_contract(rest_client, clear_chats): "name too long", {"name": "a" * (CHAT_ASSISTANT_NAME_LIMIT + 1)}, 102, - f"Chat name length is {CHAT_ASSISTANT_NAME_LIMIT + 1} which is larger than {CHAT_ASSISTANT_NAME_LIMIT}.", + f"chat name length is {CHAT_ASSISTANT_NAME_LIMIT + 1} which is larger than {CHAT_ASSISTANT_NAME_LIMIT}", None, ), - ("name wrong type", {"name": 1}, 102, "Chat name must be a string.", None), - ("name empty", {"name": ""}, 102, "`name` cannot be empty.", None), - ("duplicate lowercase", {"name": "restful_chat_update_duplicate"}, 102, "Duplicated chat name.", None), - ("duplicate uppercase", {"name": "RESTFUL_CHAT_UPDATE_DUPLICATE"}, 102, "Duplicated chat name.", None), + ("name wrong type", {"name": 1}, 102, "chat name must be a string", None), + ("name empty", {"name": ""}, 102, "`name` cannot be empty", None), + ("duplicate lowercase", {"name": "restful_chat_update_duplicate"}, 102, "duplicated chat name", None), + ("duplicate uppercase", {"name": "RESTFUL_CHAT_UPDATE_DUPLICATE"}, 102, "duplicated chat name", None), ] for scenario_name, patch_payload, expected_code, expected_message, expected_name in cases: @@ -1980,13 +1980,13 @@ def test_chat_update_mapping_and_validation_branches_p2(rest_client, clear_chats assert empty_name_res.status_code == 200 empty_name_payload = empty_name_res.json() assert empty_name_payload["code"] == 102, empty_name_payload - assert empty_name_payload["message"] == "`name` cannot be empty.", empty_name_payload + assert empty_name_payload["message"] == "`name` cannot be empty", empty_name_payload duplicate_name_res = rest_client.patch(f"/chats/{chat_id}", json={"name": "restful_chat_update_mapping_duplicate"}) assert duplicate_name_res.status_code == 200 duplicate_name_payload = duplicate_name_res.json() assert duplicate_name_payload["code"] == 102, duplicate_name_payload - assert duplicate_name_payload["message"] == "Duplicated chat name.", duplicate_name_payload + assert duplicate_name_payload["message"] == "duplicated chat name", duplicate_name_payload prompt_without_placeholder_res = rest_client.patch( f"/chats/{chat_id}", diff --git a/test/testcases/restful_api/test_datasets.py b/test/testcases/restful_api/test_datasets.py index e08dfdc3b6..a97338a76b 100644 --- a/test/testcases/restful_api/test_datasets.py +++ b/test/testcases/restful_api/test_datasets.py @@ -32,7 +32,7 @@ PARSER_ID_FIELD = "parser_id" if IS_GO_PROXY else "chunk_method" def _skip_go_ignored_null(payload, field): - if IS_GO_PROXY and payload.get("message") == "No properties were modified": + if IS_GO_PROXY and payload.get("message") == "no properties were modified": pytest.skip(f"Go dataset update ignores an explicit null {field}") @@ -676,7 +676,7 @@ def test_dataset_update_content_type_and_payload_contract(rest_client, clear_dat assert empty_payload_res.status_code == 200 empty_payload = empty_payload_res.json() assert empty_payload["code"] == 102, empty_payload - assert empty_payload["message"] == "No properties were modified", empty_payload + assert empty_payload["message"] == "no properties were modified", empty_payload unset_payload_res = rest_client.put(f"/datasets/{dataset_id}") assert unset_payload_res.status_code == 200 @@ -778,7 +778,7 @@ def test_dataset_update_description_validation_contract(rest_client, clear_datas none_res = rest_client.put(f"/datasets/{dataset_id}", json={"description": None}) assert none_res.status_code == 200 none_payload = none_res.json() - if IS_GO_PROXY and none_payload.get("message") == "No properties were modified": + if IS_GO_PROXY and none_payload.get("message") == "no properties were modified": pytest.skip("Go dataset update does not clear description with an explicit null") assert none_payload["code"] == 0, none_payload diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py index 6670e59b91..b07f5f2317 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py @@ -513,7 +513,7 @@ def test_create_chat_blank_name_is_treated_as_missing(monkeypatch): res = _run(module.create.__wrapped__()) assert res["code"] == 102 - assert res["message"] == "`name` is required." + assert res["message"] == "`name` is required" @pytest.mark.p1 diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_create_chat_assistant.py b/test/testcases/test_http_api/test_chat_assistant_management/test_create_chat_assistant.py index 4409acc5cb..b49e8caf34 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_create_chat_assistant.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_create_chat_assistant.py @@ -46,9 +46,9 @@ class TestChatAssistantCreate: ({"name": "valid_name"}, 0, ""), pytest.param({"name": "a" * (CHAT_ASSISTANT_NAME_LIMIT + 1)}, 102, "", marks=pytest.mark.skip(reason="issues/")), pytest.param({"name": 1}, 100, "", marks=pytest.mark.skip(reason="issues/")), - ({"name": ""}, 102, "`name` is required."), - ({"name": "duplicated_name"}, 102, "Duplicated chat name in creating chat."), - ({"name": "case insensitive"}, 102, "Duplicated chat name in creating chat."), + ({"name": ""}, 102, "`name` is required"), + ({"name": "duplicated_name"}, 102, "duplicated chat name in creating chat"), + ({"name": "case insensitive"}, 102, "duplicated chat name in creating chat"), ], ) def test_name(self, HttpApiAuth, add_chunks, payload, expected_code, expected_message): @@ -270,7 +270,7 @@ class TestChatAssistantCreate: tenant_payload = {"name": "guard-tenant-id", "dataset_ids": [], "tenant_id": "tenant-should-not-pass"} res = create_chat_assistant(HttpApiAuth, tenant_payload) assert res["code"] == 102 - assert res["message"] == "`tenant_id` must not be provided." + assert res["message"] == "`tenant_id` must not be provided" rerank_payload = { "name": "guard-rerank-id", diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_update_chat_assistant.py b/test/testcases/test_http_api/test_chat_assistant_management/test_update_chat_assistant.py index 08dd421e92..0c4b243fe9 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_update_chat_assistant.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_update_chat_assistant.py @@ -43,9 +43,9 @@ class TestChatAssistantUpdate: pytest.param({"name": "valid_name"}, 0, "", marks=pytest.mark.p1), pytest.param({"name": "a" * (CHAT_ASSISTANT_NAME_LIMIT + 1)}, 102, "", marks=pytest.mark.skip(reason="issues/")), pytest.param({"name": 1}, 100, "", marks=pytest.mark.skip(reason="issues/")), - pytest.param({"name": ""}, 102, "`name` cannot be empty.", marks=pytest.mark.p3), - pytest.param({"name": "test_chat_assistant_1"}, 102, "Duplicated chat name.", marks=pytest.mark.p3), - pytest.param({"name": "TEST_CHAT_ASSISTANT_1"}, 102, "Duplicated chat name.", marks=pytest.mark.p3), + pytest.param({"name": ""}, 102, "`name` cannot be empty", marks=pytest.mark.p3), + pytest.param({"name": "test_chat_assistant_1"}, 102, "duplicated chat name", marks=pytest.mark.p3), + pytest.param({"name": "TEST_CHAT_ASSISTANT_1"}, 102, "duplicated chat name", marks=pytest.mark.p3), ], ) def test_name(self, HttpApiAuth, add_chat_assistants_func, payload, expected_code, expected_message): @@ -257,12 +257,12 @@ class TestChatAssistantUpdate: # PATCH: empty name res = patch_chat_assistant(HttpApiAuth, chat_id, {"name": ""}) assert res["code"] == 102 - assert res["message"] == "`name` cannot be empty." + assert res["message"] == "`name` cannot be empty" # PATCH: duplicate name res = patch_chat_assistant(HttpApiAuth, chat_id, {"name": "test_chat_assistant_1"}) assert res["code"] == 102 - assert res["message"] == "Duplicated chat name." + assert res["message"] == "duplicated chat name" # PATCH: prompt_config without placeholder is allowed res = patch_chat_assistant( diff --git a/test/testcases/test_http_api/test_dataset_management/test_update_dataset.py b/test/testcases/test_http_api/test_dataset_management/test_update_dataset.py index d6317ea5da..c817f7533f 100644 --- a/test/testcases/test_http_api/test_dataset_management/test_update_dataset.py +++ b/test/testcases/test_http_api/test_dataset_management/test_update_dataset.py @@ -77,7 +77,7 @@ class TestRquest: dataset_id = add_dataset_func res = update_dataset(HttpApiAuth, dataset_id, {}) assert res["code"] == 102, res - assert res["message"] == "No properties were modified", res + assert res["message"] == "no properties were modified", res @pytest.mark.p3 def test_payload_unset(self, HttpApiAuth, add_dataset_func): diff --git a/test/testcases/test_sdk_api/test_chat_assistant_management/test_create_chat_assistant.py b/test/testcases/test_sdk_api/test_chat_assistant_management/test_create_chat_assistant.py index fb51252607..8d226acdcf 100644 --- a/test/testcases/test_sdk_api/test_chat_assistant_management/test_create_chat_assistant.py +++ b/test/testcases/test_sdk_api/test_chat_assistant_management/test_create_chat_assistant.py @@ -31,9 +31,9 @@ class TestChatAssistantCreate: ("valid_name", ""), pytest.param("a" * (CHAT_ASSISTANT_NAME_LIMIT + 1), "", marks=pytest.mark.skip(reason="issues/")), pytest.param(1, "", marks=pytest.mark.skip(reason="issues/")), - ("", "`name` is required."), - ("duplicated_name", "Duplicated chat name in creating chat."), - ("case insensitive", "Duplicated chat name in creating chat."), + ("", "`name` is required"), + ("duplicated_name", "duplicated chat name in creating chat"), + ("case insensitive", "duplicated chat name in creating chat"), ], ) def test_name(self, client, name, expected_message): diff --git a/test/testcases/test_sdk_api/test_chat_assistant_management/test_update_chat_assistant.py b/test/testcases/test_sdk_api/test_chat_assistant_management/test_update_chat_assistant.py index b9645518e8..5ebaf22ffb 100644 --- a/test/testcases/test_sdk_api/test_chat_assistant_management/test_update_chat_assistant.py +++ b/test/testcases/test_sdk_api/test_chat_assistant_management/test_update_chat_assistant.py @@ -75,9 +75,9 @@ class TestChatAssistantUpdate: pytest.param({"name": "valid_name"}, "", marks=pytest.mark.p1), pytest.param({"name": "a" * (CHAT_ASSISTANT_NAME_LIMIT + 1)}, "", marks=pytest.mark.skip(reason="issues/")), pytest.param({"name": 1}, "", marks=pytest.mark.skip(reason="issues/")), - pytest.param({"name": ""}, "`name` cannot be empty.", marks=pytest.mark.p3), - pytest.param({"name": "test_chat_assistant_1"}, "Duplicated chat name.", marks=pytest.mark.p3), - pytest.param({"name": "TEST_CHAT_ASSISTANT_1"}, "Duplicated chat name.", marks=pytest.mark.p3), + pytest.param({"name": ""}, "`name` cannot be empty", marks=pytest.mark.p3), + pytest.param({"name": "test_chat_assistant_1"}, "duplicated chat name", marks=pytest.mark.p3), + pytest.param({"name": "TEST_CHAT_ASSISTANT_1"}, "duplicated chat name", marks=pytest.mark.p3), ], ) def test_name(self, client, add_chat_assistants_func, payload, expected_message): diff --git a/test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py b/test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py index ca37e3c9f6..2a5b2c1d42 100644 --- a/test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py +++ b/test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py @@ -34,7 +34,7 @@ class TestRquest: dataset = add_dataset_func with pytest.raises(Exception) as exception_info: dataset.update({}) - assert "No properties were modified" in str(exception_info.value), str(exception_info.value) + assert "no properties were modified" in str(exception_info.value), str(exception_info.value) class TestCapability: