diff --git a/.gitignore b/.gitignore index fbf80b3aa..11aa54493 100644 --- a/.gitignore +++ b/.gitignore @@ -195,3 +195,6 @@ ragflow_cli.egg-info # Default backup dir backup + + +.hypothesis \ No newline at end of file diff --git a/api/apps/memory_app.py b/api/apps/memory_app.py index 363215c31..fe6caa33f 100644 --- a/api/apps/memory_app.py +++ b/api/apps/memory_app.py @@ -14,7 +14,6 @@ # limitations under the License. # import logging -import json from api.apps import login_required, current_user from api.db import TenantPermission @@ -72,7 +71,7 @@ async def update_memory(memory_id): req = await request_json() update_dict = {} # check name length - if req.get("name"): + if "name" in req: name = req["name"] memory_name = name.strip() if len(memory_name) == 0: @@ -112,7 +111,7 @@ async def update_memory(memory_id): return get_json_result(code=RetCode.NOT_FOUND, message=f"Memory '{memory_id}' not found.") memory_dict = current_memory.to_dict() - memory_dict.update({"memory_type": json.loads(current_memory.memory_type)}) + memory_dict.update({"memory_type": get_memory_type_human(current_memory.memory_type)}) to_update = {} for k, v in update_dict.items(): if isinstance(v, list) and set(memory_dict[k]) != set(v): @@ -168,3 +167,12 @@ async def list_memory(): except Exception as e: logging.error(e) return get_json_result(message=str(e), code=RetCode.SERVER_ERROR) + + +@manager.route("/config/", methods=["GET"]) # noqa: F821 +@login_required +async def get_memory_config(memory_id): + memory = MemoryService.get_with_owner_name_by_id(memory_id) + if not memory: + return get_json_result(code=RetCode.NOT_FOUND, message=f"Memory '{memory_id}' not found.") + return get_json_result(message=True, data=format_ret_data_from_memory(memory)) diff --git a/api/db/services/memory_service.py b/api/db/services/memory_service.py index 22cc6d17c..8436bf689 100644 --- a/api/db/services/memory_service.py +++ b/api/db/services/memory_service.py @@ -15,8 +15,7 @@ # from typing import List -import json - +from api.apps import current_user from api.db.db_models import DB, Memory, User from api.db.services import duplicate_name from api.db.services.common_service import CommonService @@ -32,12 +31,12 @@ class MemoryService(CommonService): @classmethod @DB.connection_context() - def get_by_memory_id(cls, memory_id: str) -> Memory: + def get_by_memory_id(cls, memory_id: str): return cls.model.select().where(cls.model.id == memory_id).first() @classmethod @DB.connection_context() - def get_by_filter(cls, filter_dict: dict, keywords: str, page: int = 1, page_size: int = 50): + def get_with_owner_name_by_id(cls, memory_id: str): fields = [ cls.model.id, cls.model.name, @@ -56,12 +55,31 @@ class MemoryService(CommonService): cls.model.system_prompt, cls.model.user_prompt ] + memory = cls.model.select(*fields).join(User, on=(cls.model.tenant_id == User.id)).where( + cls.model.id == memory_id + ).first() + return memory + + @classmethod + @DB.connection_context() + def get_by_filter(cls, filter_dict: dict, keywords: str, page: int = 1, page_size: int = 50): + fields = [ + cls.model.id, + cls.model.name, + cls.model.avatar, + cls.model.tenant_id, + User.nickname.alias("owner_name"), + cls.model.memory_type, + cls.model.storage_type, + cls.model.permissions, + cls.model.description + ] memories = cls.model.select(*fields).join(User, on=(cls.model.tenant_id == User.id)) if filter_dict.get("tenant_id"): memories = memories.where(cls.model.tenant_id.in_(filter_dict["tenant_id"])) if filter_dict.get("memory_type"): memory_type_int = calculate_memory_type(filter_dict["memory_type"]) - memories = memories.where((cls.model.memory_type & memory_type_int)) + memories = memories.where(cls.model.memory_type.bin_and(memory_type_int) > 0) if filter_dict.get("storage_type"): memories = memories.where(cls.model.storage_type == filter_dict["storage_type"]) if keywords: @@ -113,6 +131,12 @@ class MemoryService(CommonService): return 0 if "temperature" in update_dict and isinstance(update_dict["temperature"], str): update_dict["temperature"] = float(update_dict["temperature"]) + if "name" in update_dict: + update_dict["name"] = duplicate_name( + cls.query, + name=update_dict["name"], + tenant_id=current_user.id + ) update_dict.update({ "update_time": current_timestamp(), "update_date": get_format_time() diff --git a/api/utils/memory_utils.py b/api/utils/memory_utils.py index b7824c530..b8b6c2397 100644 --- a/api/utils/memory_utils.py +++ b/api/utils/memory_utils.py @@ -22,6 +22,7 @@ def format_ret_data_from_memory(memory): "name": memory.name, "avatar": memory.avatar, "tenant_id": memory.tenant_id, + "owner_name": memory.owner_name if hasattr(memory, "owner_name") else None, "memory_type": get_memory_type_human(memory.memory_type), "storage_type": memory.storage_type, "embd_id": memory.embd_id, diff --git a/test/testcases/configs.py b/test/testcases/configs.py index a94a627b7..6281645bb 100644 --- a/test/testcases/configs.py +++ b/test/testcases/configs.py @@ -19,7 +19,7 @@ import pytest HOST_ADDRESS = os.getenv("HOST_ADDRESS", "http://127.0.0.1:9380") VERSION = "v1" -ZHIPU_AI_API_KEY = os.getenv("ZHIPU_AI_API_KEY") +ZHIPU_AI_API_KEY = os.getenv("ZHIPU_AI_API_KEY", "b2e132cfedbf4f0882f66379bdd0febb.7xuKRthJvjFJm8OA") if ZHIPU_AI_API_KEY is None: pytest.exit("Error: Environment variable ZHIPU_AI_API_KEY must be set") diff --git a/test/testcases/test_web_api/common.py b/test/testcases/test_web_api/common.py index 8bd1bec1d..e1595439a 100644 --- a/test/testcases/test_web_api/common.py +++ b/test/testcases/test_web_api/common.py @@ -268,7 +268,7 @@ def create_memory(auth, payload=None): def update_memory(auth, memory_id:str, payload=None): - url = f"{HOST_ADDRESS}{MEMORY_API_URL}/{memory_id}/update" + url = f"{HOST_ADDRESS}{MEMORY_API_URL}/update/{memory_id}" res = requests.put(url=url, headers=HEADERS, auth=auth, json=payload) return res.json() @@ -283,3 +283,9 @@ def list_memory(auth, payload=None): url = f"{HOST_ADDRESS}{MEMORY_API_URL}/list" res = requests.post(url=url, headers=HEADERS, auth=auth, json=payload) return res.json() + + +def get_memory_config(auth, memory_id:str): + url = f"{HOST_ADDRESS}{MEMORY_API_URL}/config/{memory_id}" + res = requests.get(url=url, headers=HEADERS, auth=auth) + return res.json() diff --git a/test/testcases/test_web_api/test_memory_app/conftest.py b/test/testcases/test_web_api/test_memory_app/conftest.py index 8d2301c4d..11c7c2a10 100644 --- a/test/testcases/test_web_api/test_memory_app/conftest.py +++ b/test/testcases/test_web_api/test_memory_app/conftest.py @@ -15,7 +15,7 @@ # import pytest import random -from common import create_memory, list_memory, delete_memory +from test_web_api.common import create_memory, list_memory, delete_memory @pytest.fixture(scope="function") def add_memory_func(request, WebApiAuth): @@ -36,6 +36,5 @@ def add_memory_func(request, WebApiAuth): "llm_id": "ZHIPU-AI@glm-4-flash" } res = create_memory(WebApiAuth, payload) - print(res) memory_ids.append(res["data"]["id"]) return memory_ids diff --git a/test/testcases/test_web_api/test_memory_app/test_create_memory.py b/test/testcases/test_web_api/test_memory_app/test_create_memory.py index 177b91dd0..d91500bc9 100644 --- a/test/testcases/test_web_api/test_memory_app/test_create_memory.py +++ b/test/testcases/test_web_api/test_memory_app/test_create_memory.py @@ -13,3 +13,94 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import random +import re + +import pytest +from test_web_api.common import create_memory +from configs import INVALID_API_TOKEN +from libs.auth import RAGFlowWebApiAuth +from hypothesis import example, given, settings +from test.testcases.utils.hypothesis_utils import valid_names + + +class TestAuthorization: + @pytest.mark.p1 + @pytest.mark.parametrize( + "invalid_auth, expected_code, expected_message", + [ + (None, 401, ""), + (RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, ""), + ], + ids=["empty_auth", "invalid_api_token"] + ) + def test_auth_invalid(self, invalid_auth, expected_code, expected_message): + res = create_memory(invalid_auth) + assert res["code"] == expected_code, res + assert res["message"] == expected_message, res + + +class TestMemoryCreate: + @pytest.mark.p1 + @given(name=valid_names()) + @example("d" * 128) + @settings(max_examples=20) + def test_name(self, WebApiAuth, name): + payload = { + "name": name, + "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), + "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5", + "llm_id": "ZHIPU-AI@glm-4-flash" + } + res = create_memory(WebApiAuth, payload) + assert res["code"] == 0, res + pattern = rf'^{name}|{name}(?:\((\d+)\))?$' + escaped_name = re.escape(res["data"]["name"]) + assert re.match(pattern, escaped_name), res + + @pytest.mark.p2 + @pytest.mark.parametrize( + "name, expected_message", + [ + ("", "Memory name cannot be empty or whitespace."), + (" ", "Memory name cannot be empty or whitespace."), + ("a" * 129, f"Memory name '{'a'*129}' exceeds limit of 128."), + ], + ids=["empty_name", "space_name", "too_long_name"], + ) + def test_name_invalid(self, WebApiAuth, name, expected_message): + payload = { + "name": name, + "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), + "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5", + "llm_id": "ZHIPU-AI@glm-4-flash" + } + res = create_memory(WebApiAuth, payload) + assert res["message"] == expected_message, res + + @pytest.mark.p2 + @given(name=valid_names()) + def test_type_invalid(self, WebApiAuth, name): + payload = { + "name": name, + "memory_type": ["something"], + "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5", + "llm_id": "ZHIPU-AI@glm-4-flash" + } + res = create_memory(WebApiAuth, payload) + assert res["message"] == f"Memory type '{ {'something'} }' is not supported.", res + + @pytest.mark.p3 + def test_name_duplicated(self, WebApiAuth): + name = "duplicated_name_test" + payload = { + "name": name, + "memory_type": ["raw"] + random.choices(["semantic", "episodic", "procedural"], k=random.randint(0, 3)), + "embd_id": "SILICONFLOW@BAAI/bge-large-zh-v1.5", + "llm_id": "ZHIPU-AI@glm-4-flash" + } + res1 = create_memory(WebApiAuth, payload) + assert res1["code"] == 0, res1 + + res2 = create_memory(WebApiAuth, payload) + assert res2["code"] == 0, res2 diff --git a/test/testcases/test_web_api/test_memory_app/test_list_memory.py b/test/testcases/test_web_api/test_memory_app/test_list_memory.py index 72d603f4b..735599325 100644 --- a/test/testcases/test_web_api/test_memory_app/test_list_memory.py +++ b/test/testcases/test_web_api/test_memory_app/test_list_memory.py @@ -13,8 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # +from concurrent.futures import ThreadPoolExecutor, as_completed + import pytest -from common import (list_memory, delete_memory) +from test_web_api.common import list_memory, get_memory_config from configs import INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth @@ -28,26 +30,82 @@ class TestAuthorization: ], ) def test_auth_invalid(self, invalid_auth, expected_code, expected_message): - res = delete_memory(invalid_auth, "some_memory_id") + res = list_memory(invalid_auth, "some_memory_id") assert res["code"] == expected_code, res assert res["message"] == expected_message, res -class TestMemoryDelete: +class TestCapability: + @pytest.mark.p3 + def test_memory_id(self, WebApiAuth): + count = 100 + with ThreadPoolExecutor(max_workers=5) as executor: + futures = [executor.submit(list_memory, WebApiAuth) 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) + +@pytest.mark.usefixtures("add_memory_func") +class TestMemoryList: @pytest.mark.p1 - def test_memory_id(self, WebApiAuth, add_memory_func): - memory_ids = add_memory_func - res = delete_memory(WebApiAuth, memory_ids[0]) + def test_params_unset(self, WebApiAuth): + res = list_memory(WebApiAuth, None) assert res["code"] == 0, res - res = list_memory(WebApiAuth) - assert len(res["data"]["total_count"]) == 2, res + @pytest.mark.p1 + def test_params_empty(self, WebApiAuth): + res = list_memory(WebApiAuth, {}) + assert res["code"] == 0, res + + @pytest.mark.p1 + @pytest.mark.parametrize( + "params, expected_page_size", + [ + ({"page": 1, "page_size": 10}, 3), + ({"page": 2, "page_size": 10}, 0), + ({"page": 1, "page_size": 2}, 2), + ({"page": 2, "page_size": 2}, 1), + ({"page": 5, "page_size": 10}, 0), + ], + ids=["normal_first_page", "beyond_max_page", "normal_last_partial_page" , "normal_middle_page", + "full_data_single_page"], + ) + def test_page(self, WebApiAuth, params, expected_page_size): + # have added 3 memories in fixture + res = list_memory(WebApiAuth, params) + assert res["code"] == 0, res + assert len(res["data"]["memory_list"]) == expected_page_size, res @pytest.mark.p2 - @pytest.mark.usefixtures("add_memory_func") - def test_id_wrong_uuid(self, WebApiAuth): - res = delete_memory(WebApiAuth, "d94a8dc02c9711f0930f7fbc369eab6d") - assert res["code"] == 404, res + def test_filter_memory_type(self, WebApiAuth): + res = list_memory(WebApiAuth, {"filter": {"memory_type": ["semantic"]}}) + assert res["code"] == 0, res + for memory in res["data"]["memory_list"]: + assert "semantic" in memory["memory_type"], res - res = list_memory(WebApiAuth) - assert len(res["data"]["memory_list"]) == 3, res \ No newline at end of file + @pytest.mark.p2 + def test_filter_storage_type(self, WebApiAuth): + res = list_memory(WebApiAuth, {"filter":{"storage_type": "table"}}) + assert res["code"] == 0, res + for memory in res["data"]["memory_list"]: + assert memory["storage_type"] == "table", res + + @pytest.mark.p2 + def test_match_keyword(self, WebApiAuth): + res = list_memory(WebApiAuth, {"keywords": "s"}) + assert res["code"] == 0, res + for memory in res["data"]["memory_list"]: + assert "s" in memory["name"], res + + @pytest.mark.p1 + def test_get_config(self, WebApiAuth): + memory_list = list_memory(WebApiAuth, {}) + assert memory_list["code"] == 0, memory_list + + memory_config = get_memory_config(WebApiAuth, memory_list["data"]["memory_list"][0]["id"]) + assert memory_config["code"] == 0, memory_config + assert memory_config["data"]["id"] == memory_list["data"]["memory_list"][0]["id"], memory_config + for field in ["name", "avatar", "tenant_id", "owner_name", "memory_type", "storage_type", + "embd_id", "llm_id", "permissions", "description", "memory_size", "forgetting_policy", + "temperature", "system_prompt", "user_prompt"]: + assert field in memory_config["data"], memory_config diff --git a/test/testcases/test_web_api/test_memory_app/test_rm_memory.py b/test/testcases/test_web_api/test_memory_app/test_rm_memory.py index f6eb94199..e6faf5d3f 100644 --- a/test/testcases/test_web_api/test_memory_app/test_rm_memory.py +++ b/test/testcases/test_web_api/test_memory_app/test_rm_memory.py @@ -14,7 +14,7 @@ # limitations under the License. # import pytest -from common import (list_memory, delete_memory) +from test_web_api.common import (list_memory, delete_memory) from configs import INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth @@ -41,7 +41,7 @@ class TestMemoryDelete: assert res["code"] == 0, res res = list_memory(WebApiAuth) - assert len(res["data"]["total_count"]) == 2, res + assert res["data"]["total_count"] == 2, res @pytest.mark.p2 @pytest.mark.usefixtures("add_memory_func") diff --git a/test/testcases/test_web_api/test_memory_app/test_update_memory.py b/test/testcases/test_web_api/test_memory_app/test_update_memory.py index 43db7e810..4def9d8b1 100644 --- a/test/testcases/test_web_api/test_memory_app/test_update_memory.py +++ b/test/testcases/test_web_api/test_memory_app/test_update_memory.py @@ -14,9 +14,10 @@ # limitations under the License. # import pytest -from common import update_memory +from test_web_api.common import update_memory from configs import INVALID_API_TOKEN from libs.auth import RAGFlowWebApiAuth +from hypothesis import HealthCheck, example, given, settings from utils import encode_avatar from utils.file_utils import create_image_file from utils.hypothesis_utils import valid_names @@ -40,9 +41,10 @@ class TestAuthorization: class TestMemoryUpdate: - @pytest.mark.p3 + @pytest.mark.p1 @given(name=valid_names()) - @example("a" * 128) + @example("f" * 128) + @settings(max_examples=20, suppress_health_check=[HealthCheck.function_scoped_fixture]) def test_name(self, WebApiAuth, add_memory_func, name): memory_ids = add_memory_func payload = {"name": name} @@ -50,3 +52,110 @@ class TestMemoryUpdate: assert res["code"] == 0, res assert res["data"]["name"] == name, res + @pytest.mark.p2 + @pytest.mark.parametrize( + "name, expected_message", + [ + ("", "Memory name cannot be empty or whitespace."), + (" ", "Memory name cannot be empty or whitespace."), + ("a" * 129, f"Memory name '{'a' * 129}' exceeds limit of 128."), + ] + ) + def test_name_invalid(self, WebApiAuth, add_memory_func, name, expected_message): + memory_ids = add_memory_func + payload = {"name": name} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 101, res + assert res["message"] == expected_message, res + + @pytest.mark.p2 + def test_duplicate_name(self, WebApiAuth, add_memory_func): + memory_ids = add_memory_func + payload = {"name": "Test_Memory"} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + + payload = {"name": "Test_Memory"} + res = update_memory(WebApiAuth, memory_ids[1], payload) + assert res["code"] == 0, res + assert res["data"]["name"] == "Test_Memory(1)", res + + @pytest.mark.p1 + def test_avatar(self, WebApiAuth, add_memory_func, tmp_path): + memory_ids = add_memory_func + fn = create_image_file(tmp_path / "ragflow_test.png") + payload = {"avatar": f"data:image/png;base64,{encode_avatar(fn)}"} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + assert res["data"]["avatar"] == f"data:image/png;base64,{encode_avatar(fn)}", res + + @pytest.mark.p1 + def test_description(self, WebApiAuth, add_memory_func): + memory_ids = add_memory_func + description = "This is a test description." + payload = {"description": description} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + assert res["data"]["description"] == description, res + + @pytest.mark.p1 + def test_llm(self, WebApiAuth, add_memory_func): + memory_ids = add_memory_func + llm_id = "ZHIPU-AI@glm-4" + payload = {"llm_id": llm_id} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + assert res["data"]["llm_id"] == llm_id, res + + @pytest.mark.p1 + @pytest.mark.parametrize( + "permission", + [ + "me", + "team" + ], + ids=["me", "team"] + ) + def test_permission(self, WebApiAuth, add_memory_func, permission): + memory_ids = add_memory_func + payload = {"permissions": permission} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + assert res["data"]["permissions"] == permission.lower().strip(), res + + + @pytest.mark.p1 + def test_memory_size(self, WebApiAuth, add_memory_func): + memory_ids = add_memory_func + memory_size = 1048576 # 1 MB + payload = {"memory_size": memory_size} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + assert res["data"]["memory_size"] == memory_size, res + + @pytest.mark.p1 + def test_temperature(self, WebApiAuth, add_memory_func): + memory_ids = add_memory_func + temperature = 0.7 + payload = {"temperature": temperature} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + assert res["data"]["temperature"] == temperature, res + + @pytest.mark.p1 + def test_system_prompt(self, WebApiAuth, add_memory_func): + memory_ids = add_memory_func + system_prompt = "This is a system prompt." + payload = {"system_prompt": system_prompt} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + assert res["data"]["system_prompt"] == system_prompt, res + + @pytest.mark.p1 + def test_user_prompt(self, WebApiAuth, add_memory_func): + memory_ids = add_memory_func + user_prompt = "This is a user prompt." + payload = {"user_prompt": user_prompt} + res = update_memory(WebApiAuth, memory_ids[0], payload) + assert res["code"] == 0, res + assert res["data"]["user_prompt"] == user_prompt, res