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

feat(memory, model): update multi-modal memory write and model list API

Browse Source 
- Adjust multi-modal memory write behavior for text and visual data
- Mask API keys in model list response to prevent exposure
- Add capability-based filtering to the model list API
This commit is contained in:
Eternity
2026-03-24 13:54:15 +08:00
parent 2ff81ba101
commit 6bba574ca6
21 changed files with 389 additions and 401 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
api/app/repositories/memory_config_repository.py
View File

@@ -408,6 +408,9 @@ class MemoryConfigRepository:
"llm_id": db_config.llm_id,
"embedding_id": db_config.embedding_id,
"rerank_id": db_config.rerank_id,
"vision_id": db_config.vision_id,
"audio_id": db_config.audio_id,
"video_id": db_config.video_id,
"enable_llm_dedup_blockwise": db_config.enable_llm_dedup_blockwise,
"enable_llm_disambiguation": db_config.enable_llm_disambiguation,
"deep_retrieval": db_config.deep_retrieval,

12
api/app/repositories/model_repository.py
View File

@@ -1,14 +1,15 @@
from sqlalchemy.orm import Session, joinedload, selectinload
from sqlalchemy import and_, or_, func, desc, select
from typing import List, Optional, Dict, Any, Tuple
import uuid
from typing import List, Optional, Dict, Any, Tuple
from sqlalchemy import and_, or_, func, desc
from sqlalchemy.orm import Session, joinedload
from app.core.logging_config import get_db_logger
from app.models.models_model import ModelConfig, ModelApiKey, ModelType, ModelBase, model_config_api_key_association
from app.schemas.model_schema import (
ModelConfigUpdate, ModelApiKeyCreate, ModelApiKeyUpdate,
ModelConfigQuery, ModelConfigQueryNew
)
from app.core.logging_config import get_db_logger
# 获取数据库专用日志器
db_logger = get_db_logger()
@@ -137,6 +138,9 @@ class ModelConfigRepository:
type_values.append(ModelType.LLM)
filters.append(ModelConfig.type.in_(type_values))
if query.capability:
filters.append(ModelConfig.capability.contains(query.capability))
if query.is_active is not None:
filters.append(ModelConfig.is_active == query.is_active)

139
api/app/repositories/neo4j/add_nodes.py
View File

@@ -1,16 +1,19 @@
from typing import List, Optional
from app.core.logging_config import get_logger
from app.core.memory.models.graph_models import DialogueNode, StatementNode, ChunkNode, MemorySummaryNode
from app.repositories.neo4j.cypher_queries import DIALOGUE_NODE_SAVE, STATEMENT_NODE_SAVE, CHUNK_NODE_SAVE, \
MEMORY_SUMMARY_NODE_SAVE, PERCEPTUAL_NODE_SAVE, PERCEPTUAL_DIALOGUE_EDGE_SAVE
MEMORY_SUMMARY_NODE_SAVE
# 使用新的仓储层
from app.repositories.neo4j.neo4j_connector import Neo4jConnector
logger = get_logger(__name__)
async def delete_all_nodes(end_user_id: str, connector: Neo4jConnector):
"""Delete all nodes in the database."""
result = await connector.execute_query(f"MATCH (n {{end_user_id: '{end_user_id}'}}) DETACH DELETE n")
print(f"All end_user_id: {end_user_id} node and edge deleted successfully")
logger.warning(f"All end_user_id: {end_user_id} node and edge deleted successfully")
return result
@@ -25,7 +28,7 @@ async def add_dialogue_nodes(dialogues: List[DialogueNode], connector: Neo4jConn
List of created node UUIDs or None if failed
"""
if not dialogues:
print("No dialogues to save")
logger.info("No dialogues to save")
return []
try:
@@ -50,11 +53,11 @@ async def add_dialogue_nodes(dialogues: List[DialogueNode], connector: Neo4jConn
)
created_uuids = [record["uuid"] for record in result]
print(f"Successfully created {len(created_uuids)} dialogue nodes: {created_uuids}")
logger.info(f"Successfully created {len(created_uuids)} dialogue nodes: {created_uuids}")
return created_uuids
except Exception as e:
print(f"Error creating dialogue nodes: {e}")
logger.info(f"Error creating dialogue nodes: {e}")
return None
@@ -69,7 +72,7 @@ async def add_statement_nodes(statements: List[StatementNode], connector: Neo4jC
List of created node UUIDs or None if failed
"""
if not statements:
print("No statements to save")
logger.info("No statements to save")
return []
try:
@@ -122,11 +125,11 @@ async def add_statement_nodes(statements: List[StatementNode], connector: Neo4jC
)
created_uuids = [record["uuid"] for record in result]
print(f"Successfully created {len(created_uuids)} statement nodes")
logger.info(f"Successfully created {len(created_uuids)} statement nodes")
return created_uuids
except Exception as e:
print(f"Error creating statement nodes: {e}")
logger.info(f"Error creating statement nodes: {e}")
return None
@@ -141,7 +144,7 @@ async def add_chunk_nodes(chunks: List[ChunkNode], connector: Neo4jConnector) ->
List of created chunk UUIDs or None if failed
"""
if not chunks:
print("No chunk nodes to add")
logger.info("No chunk nodes to add")
return []
try:
@@ -174,16 +177,18 @@ async def add_chunk_nodes(chunks: List[ChunkNode], connector: Neo4jConnector) ->
)
created_uuids = [record["uuid"] for record in result]
print(f"Successfully created {len(created_uuids)} chunk nodes")
logger.info(f"Successfully created {len(created_uuids)} chunk nodes")
return created_uuids
except Exception as e:
print(f"Error creating chunk nodes: {e}")
logger.info(f"Error creating chunk nodes: {e}")
return None
async def add_memory_summary_nodes(summaries: List[MemorySummaryNode], connector: Neo4jConnector) -> Optional[
List[str]]:
async def add_memory_summary_nodes(
summaries: List[MemorySummaryNode],
connector: Neo4jConnector
) -> Optional[List[str]]:
"""Add memory summary nodes to Neo4j in batch.
Args:
@@ -194,7 +199,7 @@ async def add_memory_summary_nodes(summaries: List[MemorySummaryNode], connector
List of created summary node ids or None if failed
"""
if not summaries:
print("No memory summary nodes to add")
logger.info("No memory summary nodes to add")
return []
try:
@@ -220,110 +225,8 @@ async def add_memory_summary_nodes(summaries: List[MemorySummaryNode], connector
summaries=flattened
)
created_ids = [record.get("uuid") for record in result]
print(f"Successfully saved {len(created_ids)} MemorySummary nodes to Neo4j")
logger.info(f"Successfully saved {len(created_ids)} MemorySummary nodes to Neo4j")
return created_ids
except Exception as e:
print(f"Failed to save MemorySummary nodes to Neo4j: {e}")
return None
async def add_perceptual_nodes(
perceptuals: list,
connector: Neo4jConnector,
embedder_client=None,
) -> Optional[List[str]]:
"""Add perceptual memory nodes to Neo4j in batch.
Args:
perceptuals: List of MemoryPerceptualModel objects from PostgreSQL
connector: Neo4j connector instance
embedder_client: Optional embedder client for generating summary embeddings
Returns:
List of created node UUIDs or None if failed
"""
if not perceptuals:
print("No perceptual nodes to add")
return []
try:
flattened = []
for p in perceptuals:
meta = p.meta_data or {}
content_meta = meta.get("content", {})
# 生成 summary embedding（如果有 embedder_client）
summary_embedding = None
if embedder_client and p.summary:
try:
summary_embedding = (await embedder_client.response([p.summary]))[0]
except Exception as emb_err:
print(f"Failed to embed perceptual summary: {emb_err}")
flattened.append({
"id": str(p.id),
"end_user_id": str(p.end_user_id),
"perceptual_type": p.perceptual_type,
"file_path": p.file_path or "",
"file_name": p.file_name or "",
"file_ext": p.file_ext or "",
"summary": p.summary or "",
"keywords": content_meta.get("keywords", []),
"topic": content_meta.get("topic", ""),
"domain": content_meta.get("domain", ""),
"created_at": p.created_time.isoformat() if p.created_time else None,
"summary_embedding": summary_embedding,
})
result = await connector.execute_query(
PERCEPTUAL_NODE_SAVE,
perceptuals=flattened,
)
created_uuids = [record.get("uuid") for record in result]
print(f"Successfully saved {len(created_uuids)} Perceptual nodes to Neo4j")
return created_uuids
except Exception as e:
print(f"Failed to save Perceptual nodes to Neo4j: {e}")
return None
async def add_perceptual_dialogue_edges(
perceptuals: list,
dialog_id: str,
connector: Neo4jConnector,
) -> Optional[List[str]]:
"""Add edges between Perceptual nodes and Dialogue nodes.
Args:
perceptuals: List of MemoryPerceptualModel objects
dialog_id: The dialogue ID (or ref_id) to link to
connector: Neo4j connector instance
Returns:
List of created edge element IDs or None if failed
"""
if not perceptuals or not dialog_id:
return []
try:
edges = []
for p in perceptuals:
edges.append({
"perceptual_id": str(p.id),
"dialog_id": dialog_id,
"end_user_id": str(p.end_user_id),
"created_at": p.created_time.isoformat() if p.created_time else None,
})
result = await connector.execute_query(
PERCEPTUAL_DIALOGUE_EDGE_SAVE,
edges=edges,
)
created_ids = [record.get("uuid") for record in result]
print(f"Successfully saved {len(created_ids)} Perceptual-Dialogue edges to Neo4j")
return created_ids
except Exception as e:
print(f"Failed to save Perceptual-Dialogue edges: {e}")
logger.info(f"Failed to save MemorySummary nodes to Neo4j: {e}")
return None

109
api/app/repositories/neo4j/cypher_queries.py
View File

@@ -1003,60 +1003,69 @@ RETURN DISTINCT
"""
Graph_Node_query = """
MATCH (n:MemorySummary)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
0 AS priority
LIMIT $limit
MATCH (n:MemorySummary)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
0 AS priority
LIMIT $limit
UNION ALL
UNION ALL
MATCH (n:Dialogue)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
1 AS priority
LIMIT 1
MATCH (n:Dialogue)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
1 AS priority
LIMIT 1
UNION ALL
UNION ALL
MATCH (n:Statement)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
1 AS priority
LIMIT $limit
MATCH (n:Statement)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
1 AS priority
LIMIT $limit
UNION ALL
UNION ALL
MATCH (n:ExtractedEntity)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
2 AS priority
LIMIT $limit
MATCH (n:ExtractedEntity)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
2 AS priority
LIMIT $limit
UNION ALL
UNION ALL
MATCH (n:Chunk)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
3 AS priority
LIMIT $limit
MATCH (n:Chunk)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
3 AS priority
LIMIT $limit
"""
UNION ALL
MATCH (n:Perceptual)
WHERE n.end_user_id = $end_user_id
RETURN
elementId(n) AS id,
labels(n) AS labels,
properties(n) AS properties,
4 AS priority
"""
# ============================================================
# Community 节点 & BELONGS_TO_COMMUNITY 边
@@ -1340,19 +1349,19 @@ SET n += {
topic: p.topic,
domain: p.domain,
created_at: p.created_at,
file_type: p.file_type,
summary_embedding: p.summary_embedding
}
RETURN n.id AS uuid
"""
# 感知记忆与对话的关联边
PERCEPTUAL_DIALOGUE_EDGE_SAVE = """
PERCEPTUAL_CHUNK_EDGE_SAVE = """
UNWIND $edges AS edge
MATCH (p:Perceptual {id: edge.perceptual_id, end_user_id: edge.end_user_id})
MATCH (d:Dialogue {end_user_id: edge.end_user_id})
WHERE d.id = edge.dialog_id OR d.ref_id = edge.dialog_id
MERGE (d)-[r:HAS_PERCEPTUAL]->(p)
SET r.end_user_id = edge.end_user_id,
MATCH (c:Chunk {id: edge.chunk_id, end_user_id: edge.end_user_id})
MERGE (c)-[r:HAS_PERCEPTUAL]->(p)
ON CREATE SET r.end_user_id = edge.end_user_id,
r.created_at = edge.created_at
RETURN elementId(r) AS uuid
"""

72
api/app/repositories/neo4j/graph_saver.py
View File

@@ -22,13 +22,18 @@ from app.core.memory.models.graph_models import (
StatementNode,
ExtractedEntityNode,
EntityEntityEdge,
PerceptualNode,
PerceptualEdge,
)
import logging
logger = logging.getLogger(__name__)
async def save_entities_and_relationships(
entity_nodes: List[ExtractedEntityNode],
entity_entity_edges: List[EntityEntityEdge],
connector: Neo4jConnector
entity_nodes: List[ExtractedEntityNode],
entity_entity_edges: List[EntityEntityEdge],
connector: Neo4jConnector
):
"""Save entities and their relationships using graph models"""
all_entities = [entity.model_dump() for entity in entity_nodes]
@@ -73,8 +78,8 @@ async def save_entities_and_relationships(
async def save_chunk_nodes(
chunk_nodes: List[ChunkNode],
connector: Neo4jConnector
chunk_nodes: List[ChunkNode],
connector: Neo4jConnector
):
"""Save chunk nodes using graph models"""
if not chunk_nodes:
@@ -89,8 +94,8 @@ async def save_chunk_nodes(
async def save_statement_chunk_edges(
statement_chunk_edges: List[StatementChunkEdge],
connector: Neo4jConnector
statement_chunk_edges: List[StatementChunkEdge],
connector: Neo4jConnector
):
"""Save statement-chunk edges using graph models"""
if not statement_chunk_edges:
@@ -118,8 +123,8 @@ async def save_statement_chunk_edges(
async def save_statement_entity_edges(
statement_entity_edges: List[StatementEntityEdge],
connector: Neo4jConnector
statement_entity_edges: List[StatementEntityEdge],
connector: Neo4jConnector
):
"""Save statement-entity edges using graph models"""
if not statement_entity_edges:
@@ -142,7 +147,7 @@ async def save_statement_entity_edges(
if all_se_edges:
try:
await connector.execute_query(
STATEMENT_ENTITY_EDGE_SAVE,
STATEMENT_ENTITY_EDGE_SAVE,
relationships=all_se_edges
)
except Exception:
@@ -154,9 +159,11 @@ async def save_dialog_and_statements_to_neo4j(
chunk_nodes: List[ChunkNode],
statement_nodes: List[StatementNode],
entity_nodes: List[ExtractedEntityNode],
perceptual_nodes: List[PerceptualNode],
entity_edges: List[EntityEntityEdge],
statement_chunk_edges: List[StatementChunkEdge],
statement_entity_edges: List[StatementEntityEdge],
perceptual_edges: List[PerceptualEdge],
connector: Neo4jConnector,
) -> bool:
"""Save dialogue nodes, chunk nodes, statement nodes, entities, and all relationships to Neo4j using graph models.
@@ -169,9 +176,11 @@ async def save_dialog_and_statements_to_neo4j(
chunk_nodes: List of ChunkNode objects to save
statement_nodes: List of StatementNode objects to save
entity_nodes: List of ExtractedEntityNode objects to save
perceptual_nodes: List of PerceptualNode objects to save
entity_edges: List of EntityEntityEdge objects to save
statement_chunk_edges: List of StatementChunkEdge objects to save
statement_entity_edges: List of StatementEntityEdge objects to save
perceptual_edges: List of PerceptualEdge objects to save
connector: Neo4j connector instance
Returns:
@@ -190,7 +199,7 @@ async def save_dialog_and_statements_to_neo4j(
result = await tx.run(DIALOGUE_NODE_SAVE, dialogues=dialogue_data)
dialogue_uuids = [record["uuid"] async for record in result]
results['dialogues'] = dialogue_uuids
print(f"Dialogues saved to Neo4j with UUIDs: {dialogue_uuids}")
logger.info(f"Dialogues saved to Neo4j with UUIDs: {dialogue_uuids}")
# 2. Save all chunk nodes in batch
if chunk_nodes:
@@ -201,6 +210,14 @@ async def save_dialog_and_statements_to_neo4j(
results['chunks'] = chunk_uuids
logger.info(f"Successfully saved {len(chunk_uuids)} chunk nodes to Neo4j")
if perceptual_nodes:
from app.repositories.neo4j.cypher_queries import PERCEPTUAL_NODE_SAVE
perceptual_data = [node.model_dump() for node in perceptual_nodes]
result = await tx.run(PERCEPTUAL_NODE_SAVE, perceptuals=perceptual_data)
perceptual_uuids = [record["uuid"] async for record in result]
results["perceptuals"] = perceptual_uuids
logger.info(f"Successfully saved {len(perceptual_uuids)} perceptual nodes to Neo4j")
# 3. Save all statement nodes in batch
if statement_nodes:
from app.repositories.neo4j.cypher_queries import STATEMENT_NODE_SAVE
@@ -281,6 +298,22 @@ async def save_dialog_and_statements_to_neo4j(
results['statement_entity_edges'] = se_uuids
logger.info(f"Successfully saved {len(se_uuids)} statement-entity edges to Neo4j")
if perceptual_edges:
from app.repositories.neo4j.cypher_queries import PERCEPTUAL_CHUNK_EDGE_SAVE
perceptual_edge_data = []
for edge in perceptual_edges:
print(edge.source, edge.target)
perceptual_edge_data.append({
"perceptual_id": edge.source,
"chunk_id": edge.target,
"end_user_id": edge.end_user_id,
"created_at": edge.created_at.isoformat() if edge.created_at else None,
})
result = await tx.run(PERCEPTUAL_CHUNK_EDGE_SAVE, edges=perceptual_edge_data)
perceptual_edges_uuids = [record["uuid"] async for record in result]
results['perceptual_chunk_edges'] = perceptual_edges_uuids
logger.info(f"Successfully saved {len(perceptual_edges_uuids)} perceptual-chunk edges to Neo4j")
return results
try:
@@ -304,9 +337,9 @@ async def save_dialog_and_statements_to_neo4j(
def schedule_clustering_after_write(
entity_nodes: List,
llm_model_id: Optional[str] = None,
embedding_model_id: Optional[str] = None,
entity_nodes: List,
llm_model_id: Optional[str] = None,
embedding_model_id: Optional[str] = None,
) -> None:
"""
写入 Neo4j 成功后，调度后台聚类任务。
@@ -325,14 +358,15 @@ def schedule_clustering_after_write(
end_user_id = entity_nodes[0].end_user_id
new_entity_ids = [e.id for e in entity_nodes]
logger.info(f"[Clustering] 准备触发聚类，实体数: {len(new_entity_ids)}, end_user_id: {end_user_id}")
asyncio.create_task(_trigger_clustering(new_entity_ids, end_user_id, llm_model_id=llm_model_id, embedding_model_id=embedding_model_id))
asyncio.create_task(_trigger_clustering(new_entity_ids, end_user_id, llm_model_id=llm_model_id,
embedding_model_id=embedding_model_id))
async def _trigger_clustering(
new_entity_ids: List[str],
end_user_id: str,
llm_model_id: Optional[str] = None,
embedding_model_id: Optional[str] = None,
new_entity_ids: List[str],
end_user_id: str,
llm_model_id: Optional[str] = None,
embedding_model_id: Optional[str] = None,
) -> None:
"""
聚类触发函数，自动判断全量初始化还是增量更新。