diff --git a/api/apps/restful_apis/search_api.py b/api/apps/restful_apis/search_api.py index 69965b2e28..3323f971f1 100644 --- a/api/apps/restful_apis/search_api.py +++ b/api/apps/restful_apis/search_api.py @@ -25,6 +25,7 @@ from api.apps import current_user, login_required from api.constants import DATASET_NAME_LIMIT from api.db.db_models import DB from api.db.services import duplicate_name +from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.search_service import SearchService from api.db.services.user_service import TenantService, UserTenantService from common.misc_utils import get_uuid @@ -220,6 +221,11 @@ async def completion(search_id): if not kb_ids: return get_data_error_result(message="`kb_ids` is required.") + # check if the kb_ids is accessible for this user + for kb_id in kb_ids: + if not KnowledgebaseService.accessible(kb_id=kb_id, user_id=uid): + return get_data_error_result(message=f"You don't own the dataset {kb_id}") + async def stream(): nonlocal req, uid, kb_ids, search_config try: diff --git a/internal/dao/kb.go b/internal/dao/kb.go index 2bae3e0417..bb6e1d6f8f 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -259,9 +259,9 @@ func (dao *KnowledgebaseDAO) GetDetail(kbID string) (*entity.KnowledgebaseDetail // 2. If user is the owner tenant, return true // 3. If permission is "me", only owner tenant can access // 4. If permission is "team", user must be a member of the tenant -func (dao *KnowledgebaseDAO) Accessible(kbID, userID string) bool { +func (dao *KnowledgebaseDAO) Accessible(datasetID, userID string) bool { var kb entity.Knowledgebase - err := DB.Where("id = ? AND status = ?", kbID, string(entity.StatusValid)).First(&kb).Error + err := DB.Where("id = ? AND status = ?", datasetID, string(entity.StatusValid)).First(&kb).Error if err != nil { return false } diff --git a/internal/dao/search.go b/internal/dao/search.go index 49dfd84242..e0d5d3032a 100644 --- a/internal/dao/search.go +++ b/internal/dao/search.go @@ -197,8 +197,8 @@ func (dao *SearchDAO) QueryByTenantIDAndID(tenantID string, searchID string) ([] // DeleteByID deletes a search by ID (soft delete by setting status to "0") // Reference: Python common_service.py::delete_by_id -func (dao *SearchDAO) DeleteByID(id string) error { - return DB.Model(&entity.Search{}).Where("id = ?", id).Update("status", "0").Error +func (dao *SearchDAO) DeleteByID(tenantID, id string) error { + return DB.Model(&entity.Search{}).Where("tenant_id = ? AND id = ?", tenantID, id).Update("status", "0").Error } // Accessible4Deletion checks if a search can be deleted by a specific user diff --git a/internal/handler/search.go b/internal/handler/search.go index 948ab947f5..8f071c64c8 100644 --- a/internal/handler/search.go +++ b/internal/handler/search.go @@ -427,7 +427,7 @@ func (h *SearchHandler) Completion(c *gin.Context) { adapter := &service.TenantStreamAdapter{LLM: h.streamLLM, TenantID: plan.UserID, ModelID: plan.ModelID} hadError := false - for delta := range h.askService.StreamWithOptions(c.Request.Context(), adapter, plan.UserID, plan.Question, plan.KBIDs, plan.Options) { + for delta := range h.askService.StreamWithOptions(c.Request.Context(), adapter, plan.UserID, plan.Question, plan.DatasetIDs, plan.Options) { switch delta.Kind { case service.AskDeltaAnswer: writer.Write(c, sseAnswer(delta.Value, nil, false)) diff --git a/internal/service/ask_service.go b/internal/service/ask_service.go index 957a7d2f09..c5c1fbe008 100644 --- a/internal/service/ask_service.go +++ b/internal/service/ask_service.go @@ -113,16 +113,16 @@ func (s *AskService) Stream(ctx context.Context, llm StreamingLLM, userID, quest // StreamWithOptions runs Stream while allowing callers such as saved Search // apps to pass search_config retrieval options through to RetrievalTest. -func (s *AskService) StreamWithOptions(ctx context.Context, llm StreamingLLM, userID, question string, kbIDs []string, opts AskStreamOptions) <-chan AskDelta { +func (s *AskService) StreamWithOptions(ctx context.Context, llm StreamingLLM, userID, question string, datasetIDs []string, opts AskStreamOptions) <-chan AskDelta { out := make(chan AskDelta, 32) go func() { defer close(out) - s.run(ctx, llm, userID, question, kbIDs, opts, out) + s.run(ctx, llm, userID, question, datasetIDs, opts, out) }() return out } -func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question string, kbIDs []string, opts AskStreamOptions, out chan<- AskDelta) { +func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question string, datasetIDs []string, opts AskStreamOptions, out chan<- AskDelta) { // Phase 1: Retrieval. topK := DefaultAskTopK if opts.TopK != nil { @@ -138,7 +138,7 @@ func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question } req := &RetrievalTestRequest{ - Datasets: common.StringSlice(kbIDs), + Datasets: common.StringSlice(datasetIDs), Question: question, DocIDs: opts.DocIDs, UseKG: opts.UseKG, diff --git a/internal/service/search.go b/internal/service/search.go index cbfa72c027..f8fdd361ca 100644 --- a/internal/service/search.go +++ b/internal/service/search.go @@ -32,6 +32,7 @@ import ( type SearchService struct { searchDAO *dao.SearchDAO userTenantDAO *dao.UserTenantDAO + datasetDAO *dao.KnowledgebaseDAO tenantService *TenantService } @@ -40,6 +41,7 @@ func NewSearchService() *SearchService { return &SearchService{ searchDAO: dao.NewSearchDAO(), userTenantDAO: dao.NewUserTenantDAO(), + datasetDAO: dao.NewKnowledgebaseDAO(), tenantService: NewTenantService(), } } @@ -326,7 +328,7 @@ func (s *SearchService) DeleteSearch(userID string, searchID string) error { // Step 2: Execute delete (same as Python SearchService.delete_by_id) // Python: cls.model.delete().where(cls.model.id == pid).execute() - if err = s.searchDAO.DeleteByID(searchID); err != nil { + if err = s.searchDAO.DeleteByID(userID, searchID); err != nil { return fmt.Errorf("failed to delete search App %s: %w", searchID, err) } @@ -346,12 +348,12 @@ func (s *SearchService) AccessibleForCompletion(userID string, searchID string) } type SearchCompletionPlan struct { - UserID string - SearchID string - Question string - KBIDs []string - ModelID string - Options AskStreamOptions + UserID string + SearchID string + Question string + DatasetIDs []string + ModelID string + Options AskStreamOptions } func (s *SearchService) PrepareCompletion(userID, searchID string, req *SearchCompletionsRequest) (*SearchCompletionPlan, common.ErrorCode, error) { @@ -376,21 +378,28 @@ func (s *SearchService) PrepareCompletion(userID, searchID string, req *SearchCo return nil, common.CodeServerError, err } if !accessible { - return nil, common.CodeAuthenticationError, fmt.Errorf("No authorization.") + return nil, common.CodeAuthenticationError, fmt.Errorf("no authorization") } searchDetail, err := s.GetDetail(searchID) if err != nil || searchDetail == nil { - return nil, common.CodeDataError, fmt.Errorf("Cannot find search %s", searchID) + return nil, common.CodeDataError, fmt.Errorf("cannot find search %s", searchID) } searchConfig := searchConfigMapFromValue(searchDetail["search_config"]) - kbIDs := stringSliceFromSearchConfig(searchConfig["kb_ids"]) - if len(kbIDs) == 0 { - kbIDs = stringSliceFromSearchConfig(req.KBIDs) + datasetIDs := stringSliceFromSearchConfig(searchConfig["kb_ids"]) + if len(datasetIDs) == 0 { + datasetIDs = stringSliceFromSearchConfig(req.KBIDs) } - if len(kbIDs) == 0 { - return nil, common.CodeDataError, fmt.Errorf("`kb_ids` is required.") + if len(datasetIDs) == 0 { + return nil, common.CodeDataError, fmt.Errorf("`kb_ids` is required") + } + + for _, datasetID := range datasetIDs { + accessible = s.datasetDAO.Accessible(datasetID, userID) + if !accessible { + return nil, common.CodeAuthenticationError, fmt.Errorf("no authorization for dataset %s", datasetID) + } } modelID, _ := stringFromSearchConfig(searchConfig["chat_id"]) @@ -399,19 +408,20 @@ func (s *SearchService) PrepareCompletion(userID, searchID string, req *SearchCo if tenantSvc == nil { tenantSvc = NewTenantService() } - defaultModel, err := tenantSvc.GetDefaultModelName(userID, entity.ModelTypeChat) + var defaultModelName string + defaultModelName, err = tenantSvc.GetDefaultModelName(userID, entity.ModelTypeChat) if err == nil { - modelID = strings.TrimSpace(defaultModel) + modelID = strings.TrimSpace(defaultModelName) } } return &SearchCompletionPlan{ - UserID: userID, - SearchID: searchID, - Question: question, - KBIDs: kbIDs, - ModelID: modelID, - Options: askOptionsFromSearchConfig(searchID, searchConfig), + UserID: userID, + SearchID: searchID, + Question: question, + DatasetIDs: datasetIDs, + ModelID: modelID, + Options: askOptionsFromSearchConfig(searchID, searchConfig), }, common.CodeSuccess, nil } diff --git a/test/testcases/restful_api/conftest.py b/test/testcases/restful_api/conftest.py index 70a48937ba..64db36fd70 100644 --- a/test/testcases/restful_api/conftest.py +++ b/test/testcases/restful_api/conftest.py @@ -247,3 +247,12 @@ def ensure_parsed_document(rest_client, create_document): return dataset_id, document_id return _ensure + + +@pytest.hookimpl(tryfirst=True, hookwrapper=True) +def pytest_runtest_protocol(item, nextitem): + import time + start = time.perf_counter() + yield + duration = time.perf_counter() - start + print(f" [{duration:.3f}s]") diff --git a/test/testcases/restful_api/test_openai_compatible.py b/test/testcases/restful_api/test_openai_compatible.py index 8b1e3fd314..0f032103d7 100644 --- a/test/testcases/restful_api/test_openai_compatible.py +++ b/test/testcases/restful_api/test_openai_compatible.py @@ -117,7 +117,7 @@ def test_openai_compatible_invalid_chat(rest_client): assert expected_message in payload["message"], payload -@pytest.mark.p2 +@pytest.mark.p3 def test_openai_compatible_nonstream_shape(rest_client, create_chat): chat_id = create_chat("restful_openai_nonstream_chat") res = rest_client.post( @@ -127,7 +127,7 @@ def test_openai_compatible_nonstream_shape(rest_client, create_chat): "messages": [{"role": "user", "content": "hello"}], "stream": False, }, - timeout=60, + timeout=120, ) assert res.status_code == 200 payload = res.json() @@ -157,7 +157,7 @@ def test_openai_compatible_defaults_to_nonstream_when_stream_is_missing(rest_cli "model": "model", "messages": [{"role": "user", "content": "hello"}], }, - timeout=60, + timeout=120, ) assert res.status_code == 200 assert "application/json" in res.headers.get("Content-Type", ""), res.headers.get("Content-Type", "") @@ -182,7 +182,7 @@ def test_openai_compatible_nonstream_with_reference_output_shape(rest_client, cr "reference_metadata": {"include": True, "fields": ["author"]}, }, }, - timeout=60, + timeout=120, ) assert res.status_code == 200 payload = res.json() diff --git a/test/testcases/restful_api/test_searches.py b/test/testcases/restful_api/test_searches.py index 7670e94a08..2c8bab49b0 100644 --- a/test/testcases/restful_api/test_searches.py +++ b/test/testcases/restful_api/test_searches.py @@ -148,11 +148,7 @@ def test_search_completion_sse_shape_when_kb_ids_provided(rest_client, search_re timeout=60, ) assert res.status_code == 200 - content_type = res.headers.get("Content-Type", "") - assert "text/event-stream" in content_type, content_type + payload = res.json() + assert payload["code"] == 102, payload + assert "You don't own the dataset nonexistent_dataset" in payload["message"], payload - events = _sse_events(res.text) - assert events, res.text - parsed = [json.loads(evt) for evt in events] - assert isinstance(parsed[0], dict), parsed - assert parsed[-1].get("data") is True, parsed[-1]