mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-19 20:16:49 +08:00
Feat/memory (#11812)
### What problem does this PR solve? Manage and display memory datasets. ### Type of change - [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
3
.gitignore
vendored
3
.gitignore
vendored
@ -195,3 +195,6 @@ ragflow_cli.egg-info
|
|||||||
|
|
||||||
# Default backup dir
|
# Default backup dir
|
||||||
backup
|
backup
|
||||||
|
|
||||||
|
|
||||||
|
.hypothesis
|
||||||
185
api/apps/memories_app.py
Normal file
185
api/apps/memories_app.py
Normal file
@ -0,0 +1,185 @@
|
|||||||
|
#
|
||||||
|
# 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 logging
|
||||||
|
|
||||||
|
from quart import request
|
||||||
|
from api.apps import login_required, current_user
|
||||||
|
from api.db import TenantPermission
|
||||||
|
from api.db.services.memory_service import MemoryService
|
||||||
|
from api.db.services.user_service import UserTenantService
|
||||||
|
from api.utils.api_utils import validate_request, get_request_json, get_error_argument_result, get_json_result, \
|
||||||
|
not_allowed_parameters
|
||||||
|
from api.utils.memory_utils import format_ret_data_from_memory, get_memory_type_human
|
||||||
|
from api.constants import MEMORY_NAME_LIMIT, MEMORY_SIZE_LIMIT
|
||||||
|
from common.constants import MemoryType, RetCode, ForgettingPolicy
|
||||||
|
|
||||||
|
|
||||||
|
@manager.route("", methods=["POST"]) # noqa: F821
|
||||||
|
@login_required
|
||||||
|
@validate_request("name", "memory_type", "embd_id", "llm_id")
|
||||||
|
async def create_memory():
|
||||||
|
req = await get_request_json()
|
||||||
|
# check name length
|
||||||
|
name = req["name"]
|
||||||
|
memory_name = name.strip()
|
||||||
|
if len(memory_name) == 0:
|
||||||
|
return get_error_argument_result("Memory name cannot be empty or whitespace.")
|
||||||
|
if len(memory_name) > MEMORY_NAME_LIMIT:
|
||||||
|
return get_error_argument_result(f"Memory name '{memory_name}' exceeds limit of {MEMORY_NAME_LIMIT}.")
|
||||||
|
# check memory_type valid
|
||||||
|
memory_type = set(req["memory_type"])
|
||||||
|
invalid_type = memory_type - {e.name.lower() for e in MemoryType}
|
||||||
|
if invalid_type:
|
||||||
|
return get_error_argument_result(f"Memory type '{invalid_type}' is not supported.")
|
||||||
|
memory_type = list(memory_type)
|
||||||
|
|
||||||
|
try:
|
||||||
|
res, memory = MemoryService.create_memory(
|
||||||
|
tenant_id=current_user.id,
|
||||||
|
name=memory_name,
|
||||||
|
memory_type=memory_type,
|
||||||
|
embd_id=req["embd_id"],
|
||||||
|
llm_id=req["llm_id"]
|
||||||
|
)
|
||||||
|
|
||||||
|
if res:
|
||||||
|
return get_json_result(message=True, data=format_ret_data_from_memory(memory))
|
||||||
|
|
||||||
|
else:
|
||||||
|
return get_json_result(message=memory, code=RetCode.SERVER_ERROR)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
return get_json_result(message=str(e), code=RetCode.SERVER_ERROR)
|
||||||
|
|
||||||
|
|
||||||
|
@manager.route("/<memory_id>", methods=["PUT"]) # noqa: F821
|
||||||
|
@login_required
|
||||||
|
@not_allowed_parameters("id", "tenant_id", "memory_type", "storage_type", "embd_id")
|
||||||
|
async def update_memory(memory_id):
|
||||||
|
req = await get_request_json()
|
||||||
|
update_dict = {}
|
||||||
|
# check name length
|
||||||
|
if "name" in req:
|
||||||
|
name = req["name"]
|
||||||
|
memory_name = name.strip()
|
||||||
|
if len(memory_name) == 0:
|
||||||
|
return get_error_argument_result("Memory name cannot be empty or whitespace.")
|
||||||
|
if len(memory_name) > MEMORY_NAME_LIMIT:
|
||||||
|
return get_error_argument_result(f"Memory name '{memory_name}' exceeds limit of {MEMORY_NAME_LIMIT}.")
|
||||||
|
update_dict["name"] = memory_name
|
||||||
|
# check permissions valid
|
||||||
|
if req.get("permissions"):
|
||||||
|
if req["permissions"] not in [e.value for e in TenantPermission]:
|
||||||
|
return get_error_argument_result(f"Unknown permission '{req['permissions']}'.")
|
||||||
|
update_dict["permissions"] = req["permissions"]
|
||||||
|
if req.get("llm_id"):
|
||||||
|
update_dict["llm_id"] = req["llm_id"]
|
||||||
|
# check memory_size valid
|
||||||
|
if req.get("memory_size"):
|
||||||
|
if not 0 < int(req["memory_size"]) <= MEMORY_SIZE_LIMIT:
|
||||||
|
return get_error_argument_result(f"Memory size should be in range (0, {MEMORY_SIZE_LIMIT}] Bytes.")
|
||||||
|
update_dict["memory_size"] = req["memory_size"]
|
||||||
|
# check forgetting_policy valid
|
||||||
|
if req.get("forgetting_policy"):
|
||||||
|
if req["forgetting_policy"] not in [e.value for e in ForgettingPolicy]:
|
||||||
|
return get_error_argument_result(f"Forgetting policy '{req['forgetting_policy']}' is not supported.")
|
||||||
|
update_dict["forgetting_policy"] = req["forgetting_policy"]
|
||||||
|
# check temperature valid
|
||||||
|
if "temperature" in req:
|
||||||
|
temperature = float(req["temperature"])
|
||||||
|
if not 0 <= temperature <= 1:
|
||||||
|
return get_error_argument_result("Temperature should be in range [0, 1].")
|
||||||
|
update_dict["temperature"] = temperature
|
||||||
|
# allow update to empty fields
|
||||||
|
for field in ["avatar", "description", "system_prompt", "user_prompt"]:
|
||||||
|
if field in req:
|
||||||
|
update_dict[field] = req[field]
|
||||||
|
current_memory = MemoryService.get_by_memory_id(memory_id)
|
||||||
|
if not current_memory:
|
||||||
|
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": 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):
|
||||||
|
to_update[k] = v
|
||||||
|
elif memory_dict[k] != v:
|
||||||
|
to_update[k] = v
|
||||||
|
|
||||||
|
if not to_update:
|
||||||
|
return get_json_result(message=True, data=memory_dict)
|
||||||
|
|
||||||
|
try:
|
||||||
|
MemoryService.update_memory(memory_id, to_update)
|
||||||
|
updated_memory = MemoryService.get_by_memory_id(memory_id)
|
||||||
|
return get_json_result(message=True, data=format_ret_data_from_memory(updated_memory))
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(e)
|
||||||
|
return get_json_result(message=str(e), code=RetCode.SERVER_ERROR)
|
||||||
|
|
||||||
|
|
||||||
|
@manager.route("/<memory_id>", methods=["DELETE"]) # noqa: F821
|
||||||
|
@login_required
|
||||||
|
async def delete_memory(memory_id):
|
||||||
|
memory = MemoryService.get_by_memory_id(memory_id)
|
||||||
|
if not memory:
|
||||||
|
return get_json_result(message=True, code=RetCode.NOT_FOUND)
|
||||||
|
try:
|
||||||
|
MemoryService.delete_memory(memory_id)
|
||||||
|
return get_json_result(message=True)
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(e)
|
||||||
|
return get_json_result(message=str(e), code=RetCode.SERVER_ERROR)
|
||||||
|
|
||||||
|
|
||||||
|
@manager.route("", methods=["GET"]) # noqa: F821
|
||||||
|
@login_required
|
||||||
|
async def list_memory():
|
||||||
|
args = request.args
|
||||||
|
try:
|
||||||
|
tenant_ids = args.getlist("tenant_id")
|
||||||
|
memory_types = args.getlist("memory_type")
|
||||||
|
storage_type = args.get("storage_type")
|
||||||
|
keywords = args.get("keywords", "")
|
||||||
|
page = int(args.get("page", 1))
|
||||||
|
page_size = int(args.get("page_size", 50))
|
||||||
|
# make filter dict
|
||||||
|
filter_dict = {"memory_type": memory_types, "storage_type": storage_type}
|
||||||
|
if not tenant_ids:
|
||||||
|
# restrict to current user's tenants
|
||||||
|
user_tenants = UserTenantService.get_user_tenant_relation_by_user_id(current_user.id)
|
||||||
|
filter_dict["tenant_id"] = [tenant["tenant_id"] for tenant in user_tenants]
|
||||||
|
else:
|
||||||
|
filter_dict["tenant_id"] = tenant_ids
|
||||||
|
|
||||||
|
memory_list, count = MemoryService.get_by_filter(filter_dict, keywords, page, page_size)
|
||||||
|
[memory.update({"memory_type": get_memory_type_human(memory["memory_type"])}) for memory in memory_list]
|
||||||
|
return get_json_result(message=True, data={"memory_list": memory_list, "total_count": count})
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(e)
|
||||||
|
return get_json_result(message=str(e), code=RetCode.SERVER_ERROR)
|
||||||
|
|
||||||
|
|
||||||
|
@manager.route("/<memory_id>/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))
|
||||||
@ -24,3 +24,5 @@ REQUEST_MAX_WAIT_SEC = 300
|
|||||||
|
|
||||||
DATASET_NAME_LIMIT = 128
|
DATASET_NAME_LIMIT = 128
|
||||||
FILE_NAME_LEN_LIMIT = 255
|
FILE_NAME_LEN_LIMIT = 255
|
||||||
|
MEMORY_NAME_LIMIT = 128
|
||||||
|
MEMORY_SIZE_LIMIT = 10*1024*1024 # Byte
|
||||||
|
|||||||
@ -1177,6 +1177,27 @@ class EvaluationResult(DataBaseModel):
|
|||||||
db_table = "evaluation_results"
|
db_table = "evaluation_results"
|
||||||
|
|
||||||
|
|
||||||
|
class Memory(DataBaseModel):
|
||||||
|
id = CharField(max_length=32, primary_key=True)
|
||||||
|
name = CharField(max_length=128, null=False, index=False, help_text="Memory name")
|
||||||
|
avatar = TextField(null=True, help_text="avatar base64 string")
|
||||||
|
tenant_id = CharField(max_length=32, null=False, index=True)
|
||||||
|
memory_type = IntegerField(null=False, default=1, index=True, help_text="Bit flags (LSB->MSB): 1=raw, 2=semantic, 4=episodic, 8=procedural. E.g., 5 enables raw + episodic.")
|
||||||
|
storage_type = CharField(max_length=32, default='table', null=False, index=True, help_text="table|graph")
|
||||||
|
embd_id = CharField(max_length=128, null=False, index=False, help_text="embedding model ID")
|
||||||
|
llm_id = CharField(max_length=128, null=False, index=False, help_text="chat model ID")
|
||||||
|
permissions = CharField(max_length=16, null=False, index=True, help_text="me|team", default="me")
|
||||||
|
description = TextField(null=True, help_text="description")
|
||||||
|
memory_size = IntegerField(default=5242880, null=False, index=False)
|
||||||
|
forgetting_policy = CharField(max_length=32, null=False, default="fifo", index=False, help_text="lru|fifo")
|
||||||
|
temperature = FloatField(default=0.5, index=False)
|
||||||
|
system_prompt = TextField(null=True, help_text="system prompt", index=False)
|
||||||
|
user_prompt = TextField(null=True, help_text="user prompt", index=False)
|
||||||
|
|
||||||
|
class Meta:
|
||||||
|
db_table = "memory"
|
||||||
|
|
||||||
|
|
||||||
def migrate_db():
|
def migrate_db():
|
||||||
logging.disable(logging.ERROR)
|
logging.disable(logging.ERROR)
|
||||||
migrator = DatabaseMigrator[settings.DATABASE_TYPE.upper()].value(DB)
|
migrator = DatabaseMigrator[settings.DATABASE_TYPE.upper()].value(DB)
|
||||||
|
|||||||
150
api/db/services/memory_service.py
Normal file
150
api/db/services/memory_service.py
Normal file
@ -0,0 +1,150 @@
|
|||||||
|
#
|
||||||
|
# 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.
|
||||||
|
#
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
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
|
||||||
|
from api.utils.memory_utils import calculate_memory_type
|
||||||
|
from api.constants import MEMORY_NAME_LIMIT
|
||||||
|
from common.misc_utils import get_uuid
|
||||||
|
from common.time_utils import get_format_time, current_timestamp
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryService(CommonService):
|
||||||
|
# Service class for manage memory operations
|
||||||
|
model = Memory
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
@DB.connection_context()
|
||||||
|
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_with_owner_name_by_id(cls, memory_id: str):
|
||||||
|
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.embd_id,
|
||||||
|
cls.model.llm_id,
|
||||||
|
cls.model.permissions,
|
||||||
|
cls.model.description,
|
||||||
|
cls.model.memory_size,
|
||||||
|
cls.model.forgetting_policy,
|
||||||
|
cls.model.temperature,
|
||||||
|
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.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:
|
||||||
|
memories = memories.where(cls.model.name.contains(keywords))
|
||||||
|
count = memories.count()
|
||||||
|
memories = memories.order_by(cls.model.update_time.desc())
|
||||||
|
memories = memories.paginate(page, page_size)
|
||||||
|
|
||||||
|
return list(memories.dicts()), count
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
@DB.connection_context()
|
||||||
|
def create_memory(cls, tenant_id: str, name: str, memory_type: List[str], embd_id: str, llm_id: str):
|
||||||
|
# Deduplicate name within tenant
|
||||||
|
memory_name = duplicate_name(
|
||||||
|
cls.query,
|
||||||
|
name=name,
|
||||||
|
tenant_id=tenant_id
|
||||||
|
)
|
||||||
|
if len(memory_name) > MEMORY_NAME_LIMIT:
|
||||||
|
return False, f"Memory name {memory_name} exceeds limit of {MEMORY_NAME_LIMIT}."
|
||||||
|
|
||||||
|
# build create dict
|
||||||
|
memory_info = {
|
||||||
|
"id": get_uuid(),
|
||||||
|
"name": memory_name,
|
||||||
|
"memory_type": calculate_memory_type(memory_type),
|
||||||
|
"tenant_id": tenant_id,
|
||||||
|
"embd_id": embd_id,
|
||||||
|
"llm_id": llm_id,
|
||||||
|
"create_time": current_timestamp(),
|
||||||
|
"create_date": get_format_time(),
|
||||||
|
"update_time": current_timestamp(),
|
||||||
|
"update_date": get_format_time(),
|
||||||
|
}
|
||||||
|
obj = cls.model(**memory_info).save(force_insert=True)
|
||||||
|
|
||||||
|
if not obj:
|
||||||
|
return False, "Could not create new memory."
|
||||||
|
|
||||||
|
db_row = cls.model.select().where(cls.model.id == memory_info["id"]).first()
|
||||||
|
|
||||||
|
return obj, db_row
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
@DB.connection_context()
|
||||||
|
def update_memory(cls, memory_id: str, update_dict: dict):
|
||||||
|
if not update_dict:
|
||||||
|
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()
|
||||||
|
})
|
||||||
|
|
||||||
|
return cls.model.update(update_dict).where(cls.model.id == memory_id).execute()
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
@DB.connection_context()
|
||||||
|
def delete_memory(cls, memory_id: str):
|
||||||
|
return cls.model.delete().where(cls.model.id == memory_id).execute()
|
||||||
54
api/utils/memory_utils.py
Normal file
54
api/utils/memory_utils.py
Normal file
@ -0,0 +1,54 @@
|
|||||||
|
#
|
||||||
|
# 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.
|
||||||
|
#
|
||||||
|
from typing import List
|
||||||
|
from common.constants import MemoryType
|
||||||
|
|
||||||
|
def format_ret_data_from_memory(memory):
|
||||||
|
return {
|
||||||
|
"id": memory.id,
|
||||||
|
"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,
|
||||||
|
"llm_id": memory.llm_id,
|
||||||
|
"permissions": memory.permissions,
|
||||||
|
"description": memory.description,
|
||||||
|
"memory_size": memory.memory_size,
|
||||||
|
"forgetting_policy": memory.forgetting_policy,
|
||||||
|
"temperature": memory.temperature,
|
||||||
|
"system_prompt": memory.system_prompt,
|
||||||
|
"user_prompt": memory.user_prompt,
|
||||||
|
"create_time": memory.create_time,
|
||||||
|
"create_date": memory.create_date,
|
||||||
|
"update_time": memory.update_time,
|
||||||
|
"update_date": memory.update_date
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def get_memory_type_human(memory_type: int) -> List[str]:
|
||||||
|
return [mem_type.name.lower() for mem_type in MemoryType if memory_type & mem_type.value]
|
||||||
|
|
||||||
|
|
||||||
|
def calculate_memory_type(memory_type_name_list: List[str]) -> int:
|
||||||
|
memory_type = 0
|
||||||
|
type_value_map = {mem_type.name.lower(): mem_type.value for mem_type in MemoryType}
|
||||||
|
for mem_type in memory_type_name_list:
|
||||||
|
if mem_type in type_value_map:
|
||||||
|
memory_type |= type_value_map[mem_type]
|
||||||
|
return memory_type
|
||||||
@ -151,6 +151,23 @@ class Storage(Enum):
|
|||||||
OPENDAL = 6
|
OPENDAL = 6
|
||||||
GCS = 7
|
GCS = 7
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryType(Enum):
|
||||||
|
RAW = 0b0001 # 1 << 0 = 1 (0b00000001)
|
||||||
|
SEMANTIC = 0b0010 # 1 << 1 = 2 (0b00000010)
|
||||||
|
EPISODIC = 0b0100 # 1 << 2 = 4 (0b00000100)
|
||||||
|
PROCEDURAL = 0b1000 # 1 << 3 = 8 (0b00001000)
|
||||||
|
|
||||||
|
|
||||||
|
class MemoryStorageType(StrEnum):
|
||||||
|
TABLE = "table"
|
||||||
|
GRAPH = "graph"
|
||||||
|
|
||||||
|
|
||||||
|
class ForgettingPolicy(StrEnum):
|
||||||
|
FIFO = "fifo"
|
||||||
|
|
||||||
|
|
||||||
# environment
|
# environment
|
||||||
# ENV_STRONG_TEST_COUNT = "STRONG_TEST_COUNT"
|
# ENV_STRONG_TEST_COUNT = "STRONG_TEST_COUNT"
|
||||||
# ENV_RAGFLOW_SECRET_KEY = "RAGFLOW_SECRET_KEY"
|
# ENV_RAGFLOW_SECRET_KEY = "RAGFLOW_SECRET_KEY"
|
||||||
|
|||||||
@ -28,6 +28,7 @@ CHUNK_API_URL = f"/{VERSION}/chunk"
|
|||||||
DIALOG_APP_URL = f"/{VERSION}/dialog"
|
DIALOG_APP_URL = f"/{VERSION}/dialog"
|
||||||
# SESSION_WITH_CHAT_ASSISTANT_API_URL = "/api/v1/chats/{chat_id}/sessions"
|
# SESSION_WITH_CHAT_ASSISTANT_API_URL = "/api/v1/chats/{chat_id}/sessions"
|
||||||
# SESSION_WITH_AGENT_API_URL = "/api/v1/agents/{agent_id}/sessions"
|
# SESSION_WITH_AGENT_API_URL = "/api/v1/agents/{agent_id}/sessions"
|
||||||
|
MEMORY_API_URL = f"/{VERSION}/memories"
|
||||||
|
|
||||||
|
|
||||||
# KB APP
|
# KB APP
|
||||||
@ -258,3 +259,43 @@ def delete_dialogs(auth):
|
|||||||
dialog_ids = [dialog["id"] for dialog in res["data"]]
|
dialog_ids = [dialog["id"] for dialog in res["data"]]
|
||||||
if dialog_ids:
|
if dialog_ids:
|
||||||
delete_dialog(auth, {"dialog_ids": dialog_ids})
|
delete_dialog(auth, {"dialog_ids": dialog_ids})
|
||||||
|
|
||||||
|
# MEMORY APP
|
||||||
|
def create_memory(auth, payload=None):
|
||||||
|
url = f"{HOST_ADDRESS}{MEMORY_API_URL}"
|
||||||
|
res = requests.post(url=url, headers=HEADERS, auth=auth, json=payload)
|
||||||
|
return res.json()
|
||||||
|
|
||||||
|
|
||||||
|
def update_memory(auth, memory_id:str, payload=None):
|
||||||
|
url = f"{HOST_ADDRESS}{MEMORY_API_URL}/{memory_id}"
|
||||||
|
res = requests.put(url=url, headers=HEADERS, auth=auth, json=payload)
|
||||||
|
return res.json()
|
||||||
|
|
||||||
|
|
||||||
|
def delete_memory(auth, memory_id:str):
|
||||||
|
url = f"{HOST_ADDRESS}{MEMORY_API_URL}/{memory_id}"
|
||||||
|
res = requests.delete(url=url, headers=HEADERS, auth=auth)
|
||||||
|
return res.json()
|
||||||
|
|
||||||
|
|
||||||
|
def list_memory(auth, params=None):
|
||||||
|
url = f"{HOST_ADDRESS}{MEMORY_API_URL}"
|
||||||
|
if params:
|
||||||
|
query_parts = []
|
||||||
|
for key, value in params.items():
|
||||||
|
if isinstance(value, list):
|
||||||
|
for item in value:
|
||||||
|
query_parts.append(f"{key}={item}")
|
||||||
|
else:
|
||||||
|
query_parts.append(f"{key}={value}")
|
||||||
|
query_string = "&".join(query_parts)
|
||||||
|
url = f"{url}?{query_string}"
|
||||||
|
res = requests.get(url=url, headers=HEADERS, auth=auth)
|
||||||
|
return res.json()
|
||||||
|
|
||||||
|
|
||||||
|
def get_memory_config(auth, memory_id:str):
|
||||||
|
url = f"{HOST_ADDRESS}{MEMORY_API_URL}/{memory_id}/config"
|
||||||
|
res = requests.get(url=url, headers=HEADERS, auth=auth)
|
||||||
|
return res.json()
|
||||||
|
|||||||
40
test/testcases/test_web_api/test_memory_app/conftest.py
Normal file
40
test/testcases/test_web_api/test_memory_app/conftest.py
Normal file
@ -0,0 +1,40 @@
|
|||||||
|
#
|
||||||
|
# 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 pytest
|
||||||
|
import random
|
||||||
|
from test_web_api.common import create_memory, list_memory, delete_memory
|
||||||
|
|
||||||
|
@pytest.fixture(scope="function")
|
||||||
|
def add_memory_func(request, WebApiAuth):
|
||||||
|
def cleanup():
|
||||||
|
memory_list_res = list_memory(WebApiAuth)
|
||||||
|
exist_memory_ids = [memory["id"] for memory in memory_list_res["data"]["memory_list"]]
|
||||||
|
for memory_id in exist_memory_ids:
|
||||||
|
delete_memory(WebApiAuth, memory_id)
|
||||||
|
|
||||||
|
request.addfinalizer(cleanup)
|
||||||
|
|
||||||
|
memory_ids = []
|
||||||
|
for i in range(3):
|
||||||
|
payload = {
|
||||||
|
"name": f"test_memory_{i}",
|
||||||
|
"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)
|
||||||
|
memory_ids.append(res["data"]["id"])
|
||||||
|
return memory_ids
|
||||||
@ -0,0 +1,106 @@
|
|||||||
|
#
|
||||||
|
# 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 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
|
||||||
118
test/testcases/test_web_api/test_memory_app/test_list_memory.py
Normal file
118
test/testcases/test_web_api/test_memory_app/test_list_memory.py
Normal file
@ -0,0 +1,118 @@
|
|||||||
|
#
|
||||||
|
# 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.
|
||||||
|
#
|
||||||
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from test_web_api.common import list_memory, get_memory_config
|
||||||
|
from configs import INVALID_API_TOKEN
|
||||||
|
from libs.auth import RAGFlowWebApiAuth
|
||||||
|
|
||||||
|
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'>"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_auth_invalid(self, invalid_auth, expected_code, expected_message):
|
||||||
|
res = list_memory(invalid_auth)
|
||||||
|
assert res["code"] == expected_code, res
|
||||||
|
assert res["message"] == expected_message, res
|
||||||
|
|
||||||
|
|
||||||
|
class TestCapability:
|
||||||
|
@pytest.mark.p3
|
||||||
|
def test_capability(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_params_unset(self, WebApiAuth):
|
||||||
|
res = list_memory(WebApiAuth, None)
|
||||||
|
assert res["code"] == 0, 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
|
||||||
|
def test_filter_memory_type(self, WebApiAuth):
|
||||||
|
res = list_memory(WebApiAuth, {"memory_type": ["semantic"]})
|
||||||
|
assert res["code"] == 0, res
|
||||||
|
for memory in res["data"]["memory_list"]:
|
||||||
|
assert "semantic" in memory["memory_type"], res
|
||||||
|
|
||||||
|
@pytest.mark.p2
|
||||||
|
def test_filter_multi_memory_type(self, WebApiAuth):
|
||||||
|
res = list_memory(WebApiAuth, {"memory_type": ["episodic", "procedural"]})
|
||||||
|
assert res["code"] == 0, res
|
||||||
|
for memory in res["data"]["memory_list"]:
|
||||||
|
assert "episodic" in memory["memory_type"] or "procedural" in memory["memory_type"], res
|
||||||
|
|
||||||
|
@pytest.mark.p2
|
||||||
|
def test_filter_storage_type(self, WebApiAuth):
|
||||||
|
res = list_memory(WebApiAuth, {"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
|
||||||
@ -0,0 +1,53 @@
|
|||||||
|
#
|
||||||
|
# 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 pytest
|
||||||
|
from test_web_api.common import (list_memory, delete_memory)
|
||||||
|
from configs import INVALID_API_TOKEN
|
||||||
|
from libs.auth import RAGFlowWebApiAuth
|
||||||
|
|
||||||
|
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'>"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_auth_invalid(self, invalid_auth, expected_code, expected_message):
|
||||||
|
res = delete_memory(invalid_auth, "some_memory_id")
|
||||||
|
assert res["code"] == expected_code, res
|
||||||
|
assert res["message"] == expected_message, res
|
||||||
|
|
||||||
|
|
||||||
|
class TestMemoryDelete:
|
||||||
|
@pytest.mark.p1
|
||||||
|
def test_memory_id(self, WebApiAuth, add_memory_func):
|
||||||
|
memory_ids = add_memory_func
|
||||||
|
res = delete_memory(WebApiAuth, memory_ids[0])
|
||||||
|
assert res["code"] == 0, res
|
||||||
|
|
||||||
|
res = list_memory(WebApiAuth)
|
||||||
|
assert res["data"]["total_count"] == 2, 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
|
||||||
|
|
||||||
|
res = list_memory(WebApiAuth)
|
||||||
|
assert len(res["data"]["memory_list"]) == 3, res
|
||||||
@ -0,0 +1,161 @@
|
|||||||
|
#
|
||||||
|
# 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 pytest
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
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 = update_memory(invalid_auth, "memory_id")
|
||||||
|
assert res["code"] == expected_code, res
|
||||||
|
assert res["message"] == expected_message, res
|
||||||
|
|
||||||
|
|
||||||
|
class TestMemoryUpdate:
|
||||||
|
|
||||||
|
@pytest.mark.p1
|
||||||
|
@given(name=valid_names())
|
||||||
|
@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}
|
||||||
|
res = update_memory(WebApiAuth, memory_ids[0], payload)
|
||||||
|
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
|
||||||
Reference in New Issue
Block a user