adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix/optimize inerface (#183)

Browse Source 
* [changes]Optimize the time consumption of the "/end_users" interface

* [fix]Optimize the time consumption of the "/hot_memory_tags" interface

* [changes]Optimize the time consumption of the "/end_users" interface

* [fix]Optimize the time consumption of the "/hot_memory_tags" interface

* [changes]Improve the code based on AI review
This commit is contained in:
乐力齐
2026-01-23 12:21:28 +08:00
committed by GitHub
parent 7870c6c33f
commit 6e18c92a13
4 changed files with 340 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

87
api/app/services/memory_dashboard_service.py
View File

@@ -53,18 +53,28 @@ def get_workspace_end_users(
workspace_id: uuid.UUID,
current_user: User
) -> List[EndUser]:
"""获取工作空间的所有宿主"""
"""获取工作空间的所有宿主（优化版本：减少数据库查询次数）"""
business_logger.info(f"获取工作空间宿主列表: workspace_id={workspace_id}, 操作者: {current_user.username}")
try:
# 查询应用（ORM）并转换为 Pydantic 模型
# 查询应用（ORM）
apps_orm = app_repository.get_apps_by_workspace_id(db, workspace_id)
apps = [AppSchema.model_validate(h) for h in apps_orm]
app_ids = [app.id for app in apps]
end_users = []
for app_id in app_ids:
end_user_orm_list = end_user_repository.get_end_users_by_app_id(db, app_id)
end_users.extend([EndUserSchema.model_validate(h) for h in end_user_orm_list])
if not apps_orm:
business_logger.info("工作空间下没有应用")
return []
# 提取所有 app_id
app_ids = [app.id for app in apps_orm]
# 批量查询所有 end_users（一次查询而非循环查询）
from app.models.end_user_model import EndUser as EndUserModel
end_users_orm = db.query(EndUserModel).filter(
EndUserModel.app_id.in_(app_ids)
).all()
# 转换为 Pydantic 模型（只在需要时转换）
end_users = [EndUserSchema.model_validate(eu) for eu in end_users_orm]
business_logger.info(f"成功获取 {len(end_users)} 个宿主记录")
return end_users
@@ -414,6 +424,67 @@ def get_current_user_total_chunk(
business_logger.error(f"获取用户总chunk数失败: end_user_id={end_user_id} - {str(e)}")
raise
def get_users_total_chunk_batch(
end_user_ids: List[str],
db: Session,
current_user: User
) -> dict:
"""
批量获取多个用户的总chunk数（性能优化版本）
Args:
end_user_ids: 用户ID列表
db: 数据库会话
current_user: 当前用户
Returns:
字典，key为end_user_id，value为chunk总数
格式: {"user_id_1": 100, "user_id_2": 50, ...}
"""
business_logger.info(f"批量获取 {len(end_user_ids)} 个用户的总chunk数, 操作者: {current_user.username}")
try:
from app.models.document_model import Document
from sqlalchemy import func, case
if not end_user_ids:
return {}
# 构造所有文件名
file_names = [f"{user_id}.txt" for user_id in end_user_ids]
# 一次查询获取所有用户的chunk总数
# 使用 GROUP BY file_name 来分组统计
results = db.query(
Document.file_name,
func.sum(Document.chunk_num).label('total_chunk')
).filter(
Document.file_name.in_(file_names)
).group_by(
Document.file_name
).all()
# 构建结果字典
chunk_map = {}
for file_name, total_chunk in results:
# 从文件名中提取 end_user_id (去掉 .txt 后缀)
user_id = file_name.replace('.txt', '')
chunk_map[user_id] = int(total_chunk or 0)
# 对于没有记录的用户，设置为0
for user_id in end_user_ids:
if user_id not in chunk_map:
chunk_map[user_id] = 0
business_logger.info(f"成功批量获取 {len(chunk_map)} 个用户的总chunk数")
return chunk_map
except Exception as e:
business_logger.error(f"批量获取用户总chunk数失败: {str(e)}")
raise
def get_rag_content(
end_user_id: str,
limit: int,

86
api/app/services/memory_storage_service.py
View File

@@ -12,7 +12,11 @@ from datetime import datetime
from typing import Any, AsyncGenerator, Dict, List, Optional
from app.core.logging_config import get_config_logger, get_logger
from app.core.memory.analytics.hot_memory_tags import get_hot_memory_tags
from app.core.memory.analytics.hot_memory_tags import (
get_hot_memory_tags,
get_raw_tags_from_db,
filter_tags_with_llm,
)
from app.core.memory.analytics.recent_activity_stats import get_recent_activity_stats
from app.models.user_model import User
from app.repositories.data_config_repository import DataConfigRepository
@@ -515,27 +519,79 @@ async def analytics_hot_memory_tags(
) -> List[Dict[str, Any]]:
"""
获取热门记忆标签，按数量排序并返回前N个
优化策略：
1. 先从所有用户收集原始标签（不调用LLM）
2. 聚合并合并相同标签的频率
3. 排序后取前N个
4. 只调用一次LLM进行筛选
"""
workspace_id = current_user.current_workspace_id
# 获取更多标签供LLM筛选（获取limit*4个标签）
raw_limit = limit * 4
from app.services.memory_dashboard_service import get_workspace_end_users
end_users = get_workspace_end_users(db, workspace_id, current_user)
# 使用 asyncio.to_thread 避免阻塞事件循环
end_users = await asyncio.to_thread(get_workspace_end_users, db, workspace_id, current_user)
tags = []
for end_user in end_users:
tag = await get_hot_memory_tags(str(end_user.id), limit=raw_limit)
if tag:
# 将每个用户的标签列表展平到总列表中
tags.extend(tag)
# 按频率降序排序（虽然数据库已经排序，但为了确保正确性再次排序）
sorted_tags = sorted(tags, key=lambda x: x[1], reverse=True)
if not end_users:
return []
# 只返回前limit个
top_tags = sorted_tags[:limit]
return [{"name": t, "frequency": f} for t, f in top_tags]
# 步骤1: 收集所有用户的原始标签（不调用LLM）
connector = Neo4jConnector()
try:
all_raw_tags = []
for end_user in end_users:
raw_tags = await get_raw_tags_from_db(
connector,
str(end_user.id),
limit=raw_limit,
by_user=False
)
if raw_tags:
all_raw_tags.extend(raw_tags)
if not all_raw_tags:
return []
# 步骤2: 聚合相同标签的频率
tag_frequency_map = {}
for tag_name, frequency in all_raw_tags:
if tag_name in tag_frequency_map:
tag_frequency_map[tag_name] += frequency
else:
tag_frequency_map[tag_name] = frequency
# 步骤3: 按频率降序排序，取前raw_limit个
sorted_tags = sorted(
tag_frequency_map.items(),
key=lambda x: x[1],
reverse=True
)[:raw_limit]
if not sorted_tags:
return []
# 步骤4: 只调用一次LLM进行筛选
tag_names = [tag for tag, _ in sorted_tags]
# 使用第一个用户的group_id来获取LLM配置
# 因为同一工作空间下的用户应该使用相同的配置
first_end_user_id = str(end_users[0].id)
filtered_tag_names = await filter_tags_with_llm(tag_names, first_end_user_id)
# 步骤5: 根据LLM筛选结果构建最终列表（保留频率）
final_tags = []
for tag, freq in sorted_tags:
if tag in filtered_tag_names:
final_tags.append((tag, freq))
# 步骤6: 只返回前limit个
top_tags = final_tags[:limit]
return [{"name": t, "frequency": f} for t, f in top_tags]
finally:
await connector.close()
async def analytics_recent_activity_stats() -> Dict[str, Any]: