Mem0: 构建具有可扩展长期记忆的生产级 AI 智能体

# 信息提取器
class InfoExtractor:
    def __init__(self, model):
        self.model = model
    
    def extract_salient_info(self, dialogue_history):
        """从对话历史中提取显著信息"""
        # 分析对话历史
        analysis = self.model.analyze_dialogue(dialogue_history)
        
        # 提取关键实体、事件和关系
        entities = self.extract_entities(analysis)
        events = self.extract_events(analysis)
        relationships = self.extract_relationships(analysis)
        
        # 整合提取的信息
        salient_info = {
            'entities': entities,
            'events': events,
            'relationships': relationships,
            'timestamp': time.time()
        }
        
        return salient_info

# 记忆整合器
class MemoryConsolidator:
    def __init__(self, graph_memory):
        self.graph_memory = graph_memory
    
    def consolidate(self, salient_info):
        """将提取的信息整合到记忆系统中"""
        # 处理实体
        for entity in salient_info['entities']:
            self.graph_memory.add_node('entity', entity)
        
        # 处理事件
        for event in salient_info['events']:
            self.graph_memory.add_node('event', event)
        
        # 处理关系
        for relationship in salient_info['relationships']:
            self.graph_memory.add_edge(
                relationship['source'],
                relationship['target'],
                relationship['type']
            )
        
        # 更新记忆时间戳
        self.graph_memory.update_timestamp()

# 基于图的记忆
class GraphMemory:
    def __init__(self):
        self.graph = {}
        self.timestamp = time.time()
    
    def add_node(self, node_type, node_data):
        """添加节点到图中"""
        node_id = self.generate_node_id(node_type, node_data)
        if node_id not in self.graph:
            self.graph[node_id] = {
                'type': node_type,
                'data': node_data,
                'edges': {},
                'timestamp': time.time()
            }
        return node_id
    
    def add_edge(self, source_id, target_id, edge_type):
        """添加边到图中"""
        if source_id in self.graph and target_id in self.graph:
            if target_id not in self.graph[source_id]['edges']:
                self.graph[source_id]['edges'][target_id] = {
                    'type': edge_type,
                    'timestamp': time.time()
                }
    
    def retrieve(self, query, top_k=5):
        """根据查询检索相关记忆"""
        # 分析查询
        query_entities = self.extract_query_entities(query)
        
        # 基于图进行检索
        relevant_nodes = []
        for entity in query_entities:
            entity_nodes = self.find_entity_nodes(entity)
            relevant_nodes.extend(entity_nodes)
        
        # 排序和过滤
        relevant_nodes = self.rank_nodes(relevant_nodes, query)
        return relevant_nodes[:top_k]

# 记忆检索器
class MemoryRetriever:
    def __init__(self, graph_memory):
        self.graph_memory = graph_memory
    
    def retrieve_relevant_memory(self, query):
        """检索与查询相关的记忆"""
        # 在图记忆中检索
        relevant_nodes = self.graph_memory.retrieve(query)
        
        # 格式化检索到的记忆
        formatted_memory = self.format_memory(relevant_nodes)
        
        return formatted_memory

# Mem0 主类
class Mem0:
    def __init__(self, llm, info_extractor, memory_consolidator, memory_retriever):
        self.llm = llm
        self.info_extractor = info_extractor
        self.memory_consolidator = memory_consolidator
        self.memory_retriever = memory_retriever
    
    def process_dialogue(self, user_input, dialogue_history):
        """处理对话"""
        # 检索相关记忆
        relevant_memory = self.memory_retriever.retrieve_relevant_memory(user_input)
        
        # 将记忆注入到上下文
        enhanced_context = self.inject_memory(user_input, relevant_memory, dialogue_history)
        
        # 生成响应
        response = self.llm.generate(enhanced_context)
        
        # 更新对话历史
        new_history = dialogue_history + [(user_input, response)]
        
        # 提取和整合新信息
        salient_info = self.info_extractor.extract_salient_info(new_history)
        self.memory_consolidator.consolidate(salient_info)
        
        return response, new_history