# # Copyright 2025 The InfiniFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import os from concurrent.futures import ThreadPoolExecutor, as_completed from time import sleep import pytest from common import add_chunk, delete_chunks, retrieval_chunks from configs import INVALID_API_TOKEN from libs.auth import RAGFlowHttpApiAuth @pytest.mark.p1 class TestAuthorization: @pytest.mark.parametrize( "invalid_auth, expected_code, expected_message", [ (None, 0, "`Authorization` can't be empty"), ( RAGFlowHttpApiAuth(INVALID_API_TOKEN), 109, "Authentication error: API key is invalid!", ), ], ) def test_invalid_auth(self, invalid_auth, expected_code, expected_message): res = retrieval_chunks(invalid_auth) assert res["code"] == expected_code assert res["message"] == expected_message class TestChunksRetrieval: @pytest.mark.p1 @pytest.mark.parametrize( "payload, expected_code, expected_page_size, expected_message", [ ({"question": "chunk", "dataset_ids": None}, 0, 4, ""), ({"question": "chunk", "document_ids": None}, 102, 0, "`dataset_ids` is required."), ({"question": "chunk", "dataset_ids": None, "document_ids": None}, 0, 4, ""), ({"question": "chunk"}, 102, 0, "`dataset_ids` is required."), ], ) def test_basic_scenarios(self, HttpApiAuth, add_chunks, payload, expected_code, expected_page_size, expected_message): dataset_id, document_id, _ = add_chunks if "dataset_ids" in payload: payload["dataset_ids"] = [dataset_id] if "document_ids" in payload: payload["document_ids"] = [document_id] res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == expected_code if expected_code == 0: assert len(res["data"]["chunks"]) == expected_page_size else: assert res["message"] == expected_message @pytest.mark.p2 @pytest.mark.parametrize( "payload, expected_code, expected_page_size, expected_message", [ pytest.param( {"page": None, "page_size": 2}, 100, 2, """TypeError("int() argument must be a string, a bytes-like object or a real number, not \'NoneType\'")""", marks=pytest.mark.skip, ), pytest.param( {"page": 0, "page_size": 2}, 100, 0, "ValueError('Search does not support negative slicing.')", marks=pytest.mark.skip, ), ({"page": 2, "page_size": 2}, 0, 2, ""), ({"page": 3, "page_size": 2}, 0, 0, ""), ({"page": "3", "page_size": 2}, 0, 0, ""), pytest.param( {"page": -1, "page_size": 2}, 100, 0, "ValueError('Search does not support negative slicing.')", marks=pytest.mark.skip, ), pytest.param( {"page": "a", "page_size": 2}, 100, 0, """ValueError("invalid literal for int() with base 10: \'a\'")""", marks=pytest.mark.skip, ), ], ) def test_page(self, HttpApiAuth, add_chunks, payload, expected_code, expected_page_size, expected_message): dataset_id, _, _ = add_chunks payload.update({"question": "chunk", "dataset_ids": [dataset_id]}) res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == expected_code if expected_code == 0: assert len(res["data"]["chunks"]) == expected_page_size else: assert res["message"] == expected_message @pytest.mark.p3 @pytest.mark.parametrize( "payload, expected_code, expected_page_size, expected_message", [ pytest.param( {"page_size": None}, 100, 0, """TypeError("int() argument must be a string, a bytes-like object or a real number, not \'NoneType\'")""", marks=pytest.mark.skip, ), # ({"page_size": 0}, 0, 0, ""), pytest.param({"page_size": 1}, 0, 1, "", marks=pytest.mark.skip(reason="issues/10692")), ({"page_size": 5}, 0, 4, ""), pytest.param({"page_size": "1"}, 0, 1, "", marks=pytest.mark.skip(reason="issues/10692")), # ({"page_size": -1}, 0, 0, ""), pytest.param( {"page_size": "a"}, 100, 0, """ValueError("invalid literal for int() with base 10: \'a\'")""", marks=pytest.mark.skip, ), ], ) def test_page_size(self, HttpApiAuth, add_chunks, payload, expected_code, expected_page_size, expected_message): dataset_id, _, _ = add_chunks payload.update({"question": "chunk", "dataset_ids": [dataset_id]}) res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == expected_code if expected_code == 0: assert len(res["data"]["chunks"]) == expected_page_size else: assert res["message"] == expected_message @pytest.mark.p3 @pytest.mark.parametrize( "payload, expected_code, expected_page_size, expected_message", [ ({"vector_similarity_weight": 0}, 0, 4, ""), ({"vector_similarity_weight": 0.5}, 0, 4, ""), ({"vector_similarity_weight": 10}, 0, 4, ""), pytest.param( {"vector_similarity_weight": "a"}, 100, 0, """ValueError("could not convert string to float: \'a\'")""", marks=pytest.mark.skip, ), ], ) def test_vector_similarity_weight(self, HttpApiAuth, add_chunks, payload, expected_code, expected_page_size, expected_message): dataset_id, _, _ = add_chunks payload.update({"question": "chunk", "dataset_ids": [dataset_id]}) res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == expected_code if expected_code == 0: assert len(res["data"]["chunks"]) == expected_page_size else: assert res["message"] == expected_message @pytest.mark.p3 @pytest.mark.parametrize( "payload, expected_code, expected_page_size, expected_message", [ ({"top_k": 10}, 0, 4, ""), pytest.param( {"top_k": 1}, 0, 4, "", marks=pytest.mark.skipif(os.getenv("DOC_ENGINE") in ["infinity", "opensearch"], reason="Infinity"), ), pytest.param( {"top_k": 1}, 0, 1, "", marks=pytest.mark.skipif(os.getenv("DOC_ENGINE") in [None, "opensearch", "elasticsearch"], reason="elasticsearch"), ), ({"top_k": -1}, 100, 4, "must be greater than 0"), pytest.param( {"top_k": "a"}, 100, 0, """ValueError("invalid literal for int() with base 10: \'a\'")""", marks=pytest.mark.skip, ), ], ) def test_top_k(self, HttpApiAuth, add_chunks, payload, expected_code, expected_page_size, expected_message): dataset_id, _, _ = add_chunks payload.update({"question": "chunk", "dataset_ids": [dataset_id]}) res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == expected_code if expected_code == 0: assert len(res["data"]["chunks"]) == expected_page_size else: assert expected_message in res["message"] @pytest.mark.skip @pytest.mark.parametrize( "payload, expected_code, expected_message", [ ({"rerank_id": "BAAI/bge-reranker-v2-m3"}, 0, ""), pytest.param({"rerank_id": "unknown"}, 100, "LookupError('Model(unknown) not authorized')", marks=pytest.mark.skip), ], ) def test_rerank_id(self, HttpApiAuth, add_chunks, payload, expected_code, expected_message): dataset_id, _, _ = add_chunks payload.update({"question": "chunk", "dataset_ids": [dataset_id]}) res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == expected_code if expected_code == 0: assert len(res["data"]["chunks"]) > 0 else: assert expected_message in res["message"] @pytest.mark.skip @pytest.mark.parametrize( "payload, expected_code, expected_page_size, expected_message", [ ({"keyword": True}, 0, 5, ""), ({"keyword": "True"}, 0, 5, ""), ({"keyword": False}, 0, 5, ""), ({"keyword": "False"}, 0, 5, ""), ({"keyword": None}, 0, 5, ""), ], ) def test_keyword(self, HttpApiAuth, add_chunks, payload, expected_code, expected_page_size, expected_message): dataset_id, _, _ = add_chunks payload.update({"question": "chunk test", "dataset_ids": [dataset_id]}) res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == expected_code if expected_code == 0: assert len(res["data"]["chunks"]) == expected_page_size else: assert res["message"] == expected_message @pytest.mark.p3 @pytest.mark.parametrize( "payload, expected_code, expected_highlight, expected_message", [ ({"highlight": True}, 0, True, ""), ({"highlight": "True"}, 0, True, ""), ({"highlight": False}, 0, False, ""), ({"highlight": "False"}, 0, False, ""), pytest.param({"highlight": None}, 0, False, "", marks=pytest.mark.skip(reason="issues/6648")), ], ) def test_highlight(self, HttpApiAuth, add_chunks, payload, expected_code, expected_highlight, expected_message): dataset_id, _, _ = add_chunks payload.update({"question": "chunk", "dataset_ids": [dataset_id]}) res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == expected_code if expected_highlight: for chunk in res["data"]["chunks"]: assert "highlight" in chunk else: for chunk in res["data"]["chunks"]: assert "highlight" not in chunk if expected_code != 0: assert res["message"] == expected_message @pytest.mark.p3 def test_invalid_params(self, HttpApiAuth, add_chunks): dataset_id, _, _ = add_chunks payload = {"question": "chunk", "dataset_ids": [dataset_id], "a": "b"} res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == 0 assert len(res["data"]["chunks"]) == 4 @pytest.mark.p3 def test_concurrent_retrieval(self, HttpApiAuth, add_chunks): dataset_id, _, _ = add_chunks count = 100 payload = {"question": "chunk", "dataset_ids": [dataset_id]} with ThreadPoolExecutor(max_workers=5) as executor: futures = [executor.submit(retrieval_chunks, HttpApiAuth, payload) for i in range(count)] responses = list(as_completed(futures)) assert len(responses) == count, responses assert all(future.result()["code"] == 0 for future in futures) class TestDeletedChunksNotRetrievable: """Regression tests for issue #12520: deleted slices should not appear in retrieval/reference.""" @pytest.mark.p1 def test_deleted_chunk_not_in_retrieval(self, HttpApiAuth, add_document): """ Test that a deleted chunk is not returned by the retrieval API. Steps: 1. Add a chunk with unique content 2. Verify the chunk is retrievable 3. Delete the chunk 4. Verify the chunk is no longer retrievable """ dataset_id, document_id = add_document # Add a chunk with unique content that we can search for unique_content = "UNIQUE_TEST_CONTENT_12520_REGRESSION" res = add_chunk(HttpApiAuth, dataset_id, document_id, {"content": unique_content}) assert res["code"] == 0, f"Failed to add chunk: {res}" chunk_id = res["data"]["chunk"]["id"] # Wait for indexing to complete sleep(2) # Verify the chunk is retrievable payload = {"question": unique_content, "dataset_ids": [dataset_id]} res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == 0, f"Retrieval failed: {res}" chunk_ids_before = [c["id"] for c in res["data"]["chunks"]] assert chunk_id in chunk_ids_before, f"Chunk {chunk_id} should be retrievable before deletion" # Delete the chunk res = delete_chunks(HttpApiAuth, dataset_id, document_id, {"chunk_ids": [chunk_id]}) assert res["code"] == 0, f"Failed to delete chunk: {res}" # Wait for deletion to propagate sleep(1) # Verify the chunk is no longer retrievable res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == 0, f"Retrieval failed after deletion: {res}" chunk_ids_after = [c["id"] for c in res["data"]["chunks"]] assert chunk_id not in chunk_ids_after, f"Chunk {chunk_id} should NOT be retrievable after deletion" @pytest.mark.p2 def test_deleted_chunks_batch_not_in_retrieval(self, HttpApiAuth, add_document): """ Test that multiple deleted chunks are not returned by retrieval. """ dataset_id, document_id = add_document # Add multiple chunks with unique content chunk_ids = [] for i in range(3): unique_content = f"BATCH_DELETE_TEST_CHUNK_{i}_12520" res = add_chunk(HttpApiAuth, dataset_id, document_id, {"content": unique_content}) assert res["code"] == 0, f"Failed to add chunk {i}: {res}" chunk_ids.append(res["data"]["chunk"]["id"]) # Wait for indexing sleep(2) # Verify chunks are retrievable payload = {"question": "BATCH_DELETE_TEST_CHUNK", "dataset_ids": [dataset_id]} res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == 0 retrieved_ids_before = [c["id"] for c in res["data"]["chunks"]] for cid in chunk_ids: assert cid in retrieved_ids_before, f"Chunk {cid} should be retrievable before deletion" # Delete all chunks res = delete_chunks(HttpApiAuth, dataset_id, document_id, {"chunk_ids": chunk_ids}) assert res["code"] == 0, f"Failed to delete chunks: {res}" # Wait for deletion to propagate sleep(1) # Verify none of the chunks are retrievable res = retrieval_chunks(HttpApiAuth, payload) assert res["code"] == 0 retrieved_ids_after = [c["id"] for c in res["data"]["chunks"]] for cid in chunk_ids: assert cid not in retrieved_ids_after, f"Chunk {cid} should NOT be retrievable after deletion"