feat(memory): add end user memory count filtering
- Sync memory_count after Neo4j write and forgetting cycle - Filter Neo4j end user list by memory_count > 0 - Filter RAG end user list by Memory knowledge chunk count
This commit is contained in:
@@ -313,6 +313,9 @@ async def write(
|
||||
except Exception as cache_err:
|
||||
logger.warning(f"[WRITE] 写入活动统计缓存失败(不影响主流程): {cache_err}", exc_info=True)
|
||||
|
||||
#同步neo4j记忆节点总数到pgsql,end_user表的memory_count字段
|
||||
await _sync_memory_count_after_write(end_user_id)
|
||||
|
||||
# Close LLM/Embedder underlying httpx clients to prevent
|
||||
# 'RuntimeError: Event loop is closed' during garbage collection
|
||||
for client_obj in (llm_client, embedder_client):
|
||||
@@ -331,3 +334,49 @@ async def write(
|
||||
|
||||
logger.info("=== Pipeline Complete ===")
|
||||
logger.info(f"Total execution time: {total_time:.2f} seconds")
|
||||
|
||||
|
||||
async def _sync_memory_count_after_write(end_user_id: str) -> None:
|
||||
"""
|
||||
记忆写入完成后,查 Neo4j 全量节点数,绝对值同步到 PostgreSQL end_user 表的 memory_count 字段
|
||||
|
||||
不使用增量累加:
|
||||
- Neo4j 写入使用 MERGE 语义,节点列表长度不等于新增节点数。
|
||||
- 重试或重复写入可能匹配已有节点。
|
||||
- 绝对值覆盖可以避免越加越大的计数漂移。
|
||||
"""
|
||||
if not end_user_id:
|
||||
return
|
||||
|
||||
try:
|
||||
from app.models.end_user_model import EndUser
|
||||
from app.repositories.memory_config_repository import MemoryConfigRepository
|
||||
|
||||
connector = Neo4jConnector()
|
||||
try:
|
||||
result = await connector.execute_query(
|
||||
MemoryConfigRepository.SEARCH_FOR_ALL_BATCH,
|
||||
end_user_ids=[end_user_id],
|
||||
)
|
||||
node_count = int(result[0]["total"]) if result else 0
|
||||
finally:
|
||||
await connector.close()
|
||||
|
||||
with get_db_context() as db:
|
||||
db.query(EndUser).filter(
|
||||
EndUser.id == uuid.UUID(end_user_id)
|
||||
).update(
|
||||
{"memory_count": node_count},
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.commit()
|
||||
|
||||
logger.info(
|
||||
f"[MemoryCount] 写入后同步 memory_count: "
|
||||
f"end_user_id={end_user_id}, count={node_count}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"[MemoryCount] 写入后同步 memory_count 失败(不影响主流程): {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
@@ -145,7 +145,8 @@ class ForgettingScheduler:
|
||||
}
|
||||
|
||||
logger.info("没有可遗忘的节点对,遗忘周期结束")
|
||||
|
||||
# 同步 Neo4j 记忆节点总数到 PostgreSQL的 end_user 表的 memory_count 字段
|
||||
await self._sync_memory_count_to_mysql(end_user_id)
|
||||
return report
|
||||
|
||||
# 步骤3:按激活值排序(激活值最低的优先)
|
||||
@@ -302,7 +303,8 @@ class ForgettingScheduler:
|
||||
f"({reduction_rate:.2%}), "
|
||||
f"耗时 {duration:.2f} 秒"
|
||||
)
|
||||
|
||||
# 同步 Neo4j 记忆节点总数到 PostgreSQL的 end_user 表的 memory_count 字段
|
||||
await self._sync_memory_count_to_mysql(end_user_id)
|
||||
return report
|
||||
|
||||
except Exception as e:
|
||||
@@ -350,3 +352,48 @@ class ForgettingScheduler:
|
||||
if results:
|
||||
return results[0]['total']
|
||||
return 0
|
||||
|
||||
async def _sync_memory_count_to_mysql(
|
||||
self,
|
||||
end_user_id: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
遗忘周期结束后,用 SEARCH_FOR_ALL_BATCH 口径查全量节点数,
|
||||
同步到 PostgreSQL end_users.memory_count。
|
||||
|
||||
不复用 _count_knowledge_nodes:
|
||||
- _count_knowledge_nodes 只统计 Statement、ExtractedEntity、MemorySummary。
|
||||
- 宿主列表需要统计该 end_user_id 下全部 Neo4j 节点。
|
||||
"""
|
||||
if not end_user_id:
|
||||
return
|
||||
|
||||
try:
|
||||
from app.db import get_db_context
|
||||
from app.models.end_user_model import EndUser
|
||||
from app.repositories.memory_config_repository import MemoryConfigRepository
|
||||
|
||||
result = await self.connector.execute_query(
|
||||
MemoryConfigRepository.SEARCH_FOR_ALL_BATCH,
|
||||
end_user_ids=[end_user_id],
|
||||
)
|
||||
node_count = int(result[0]["total"]) if result else 0
|
||||
|
||||
with get_db_context() as db:
|
||||
db.query(EndUser).filter(
|
||||
EndUser.id == UUID(end_user_id)
|
||||
).update(
|
||||
{"memory_count": node_count},
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.commit()
|
||||
|
||||
logger.info(
|
||||
f"[MemoryCount] 遗忘后同步 memory_count: "
|
||||
f"end_user_id={end_user_id}, count={node_count}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"[MemoryCount] 遗忘后同步 memory_count 失败(不影响主流程): {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user