feat: memory

.gitignore

@@ -195,3 +195,6 @@ ragflow_cli.egg-info
 
 # Default backup dir
 backup
+
+
+.hypothesis

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/<memory_id>", 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))

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()

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,

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")
 

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()

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

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, "<Unauthorized '401: Unauthorized'>"),
+            (RAGFlowWebApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
+        ],
+        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

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):
+    def test_params_unset(self, WebApiAuth):
-        memory_ids = add_memory_func
+        res  = list_memory(WebApiAuth, None)
-        res = delete_memory(WebApiAuth, memory_ids[0])
         assert res["code"] == 0, res
 
-        res = list_memory(WebApiAuth)
+    @pytest.mark.p1
-        assert len(res["data"]["total_count"]) == 2, res
+    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_filter_memory_type(self, WebApiAuth):
-    def test_id_wrong_uuid(self, WebApiAuth):
+        res = list_memory(WebApiAuth, {"filter": {"memory_type": ["semantic"]}})
-        res = delete_memory(WebApiAuth, "d94a8dc02c9711f0930f7fbc369eab6d")
+        assert res["code"] == 0, res
-        assert res["code"] == 404, res
+        for memory in res["data"]["memory_list"]:
+            assert "semantic" in memory["memory_type"], res
 
-        res = list_memory(WebApiAuth)
+    @pytest.mark.p2
-        assert len(res["data"]["memory_list"]) == 3, res
+    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

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")

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