AI Agent 记忆系统深度解析¶

import chromadb
from sentence_transformers import SentenceTransformer

# 初始化
client = chromadb.Client()
collection = client.create_collection("agent_memory")

# 存储记忆
def store_memory(text, metadata=None):
    collection.add(
        documents=[text],
        metadatas=[metadata or {}],
        ids=[str(hash(text))]
    )

# 检索记忆
def retrieve_memory(query, n_results=5):
    results = collection.query(
        query_texts=[query],
        n_results=n_results
    )
    return results['documents'][0]

# 示例
store_memory("用户偏好使用 Python 3.11", {"type": "preference"})
store_memory("项目使用 PostgreSQL 数据库", {"type": "project"})

results = retrieve_memory("用户喜欢什么编程语言？")
print(results)  # 输出: ["用户偏好使用 Python 3.11"]

from openai import OpenAI

client = OpenAI()

def encode_memory(text):
    """将文本编码为向量"""
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return response.data[0].embedding

def extract_entities(text):
    """使用 LLM 提取实体"""
    prompt = f"""
    从以下文本中提取关键实体和关系：
    {text}

    返回 JSON 格式：
    {{
        "entities": [
            {{"name": "实体名", "type": "类型", "description": "描述"}}
        ],
        "relations": [
            {{"source": "源实体", "target": "目标实体", "relation": "关系类型"}}
        ]
    }}
    """
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}],
        response_format={"type": "json_object"}
    )
    return json.loads(response.choices[0].message.content)

def hybrid_retrieve(query, user_id, top_k=5):
    """混合检索策略"""
    # 1. 语义检索
    semantic_results = vector_search(query, top_k=top_k*2)

    # 2. 关键词检索
    keyword_results = full_text_search(query, top_k=top_k*2)

    # 3. 时间加权
    time_weighted_results = apply_time_decay(
        semantic_results + keyword_results,
        decay_rate=0.1  # 每天衰减 10%
    )

    # 4. 去重排序
    final_results = deduplicate_and_rank(
        time_weighted_results,
        top_k=top_k
    )

    return final_results

def apply_time_decay(results, decay_rate):
    """应用时间衰减"""
    import time
    current_time = time.time()

    for result in results:
        age_days = (current_time - result['timestamp']) / 86400
        result['score'] *= (1 - decay_rate) ** age_days

    return results

def update_memory(memory_id, new_info):
    """更新记忆"""
    # 1. 获取旧记忆
    old_memory = get_memory(memory_id)

    # 2. 判断是否需要更新
    if should_update(old_memory, new_info):
        # 3. 合并信息
        merged = merge_memory(old_memory, new_info)

        # 4. 重新编码
        new_embedding = encode_memory(merged['content'])

        # 5. 更新存储
        update_in_db(memory_id, merged, new_embedding)

        # 6. 记录更新历史
        log_update(memory_id, old_memory, merged)

def should_update(old, new):
    """判断是否需要更新"""
    # 时间衰减
    if is_outdated(old, max_age_days=30):
        return True

    # 信息冲突
    if conflicts(old, new):
        return True

    # 信息补充
    if is_more_complete(old, new):
        return True

    return False

def forget_strategy(memory):
    """遗忘策略"""
    # 1. 时间衰减
    if memory.age_days > 90:
        return "archive"

    # 2. 访问频率
    if memory.access_count < 3 and memory.age_days > 30:
        return "delete"

    # 3. 重要性评估
    if memory.importance_score < 0.3:
        return "delete"

    # 4. 冲突检测
    if has_conflicts(memory):
        return "resolve_conflict"

    return "keep"

from langchain.memory import (
    ConversationBufferMemory,      # 完整对话历史
    ConversationSummaryMemory,     # 对话摘要
    ConversationBufferWindowMemory, # 滑动窗口
    ConversationKGMemory,          # 知识图谱记忆
    VectorStoreRetrieverMemory,    # 向量存储记忆
)

# 1. 缓冲记忆（最简单）
buffer_memory = ConversationBufferMemory(
    return_messages=True,
    memory_key="history"
)

# 2. 摘要记忆（适合长对话）
summary_memory = ConversationSummaryMemory(
    llm=ChatOpenAI(),
    memory_key="history"
)

# 3. 窗口记忆（平衡成本和效果）
window_memory = ConversationBufferWindowMemory(
    k=10,  # 保留最近 10 轮对话
    return_messages=True,
    memory_key="history"
)

# 4. 知识图谱记忆（适合复杂关系）
kg_memory = ConversationKGMemory(
    llm=ChatOpenAI(),
    memory_key="history"
)

# 5. 向量存储记忆（语义检索）
vector_memory = VectorStoreRetrieverMemory(
    retriever=vectorstore.as_retriever(),
    memory_key="history"
)

# HNSW 索引（推荐）
collection = client.create_collection(
    name="memories",
    metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
)

# 批量插入
collection.add(
    documents=[...],
    embeddings=[...],
    ids=[...]
)

# 创建索引
collection.create_index()

from functools import lru_cache
import hashlib

@lru_cache(maxsize=1000)
def cached_retrieve(query_hash, top_k=5):
    """缓存检索结果"""
    return retrieve_memory(query_hash, top_k)

def retrieve_with_cache(query, top_k=5):
    """带缓存的检索"""
    query_hash = hashlib.md5(query.encode()).hexdigest()
    return cached_retrieve(query_hash, top_k)

def batch_store(memories, batch_size=100):
    """批量存储记忆"""
    for i in range(0, len(memories), batch_size):
        batch = memories[i:i+batch_size]
        collection.add(
            documents=[m['content'] for m in batch],
            metadatas=[m['metadata'] for m in batch],
            ids=[m['id'] for m in batch]
        )

# 最小可行记忆系统
class SimpleMemory:
    def __init__(self):
        self.short_term = []  # 当前会话
        self.long_term = {}   # 持久化存储

    def remember(self, key, value):
        self.long_term[key] = value

    def recall(self, key):
        return self.long_term.get(key)

    def forget(self, key):
        if key in self.long_term:
            del self.long_term[key]

# 完整记忆系统
class AdvancedMemory:
    def __init__(self):
        self.vector_store = ChromaDB()
        self.graph_store = Neo4j()
        self.cache = Redis()

    def store(self, memory):
        # 1. 编码
        embedding = encode(memory.content)

        # 2. 存储到向量数据库
        self.vector_store.add(memory, embedding)

        # 3. 提取实体和关系
        entities, relations = extract(memory.content)

        # 4. 存储到知识图谱
        self.graph_store.add(entities, relations)

        # 5. 更新缓存
        self.cache.set(memory.id, memory)

    def retrieve(self, query, context):
        # 1. 缓存检查
        cached = self.cache.get(query)
        if cached:
            return cached

        # 2. 向量检索
        vector_results = self.vector_store.search(query)

        # 3. 图谱检索
        graph_results = self.graph_store.query(query)

        # 4. 融合排序
        results = merge_and_rank(vector_results, graph_results)

        # 5. 更新缓存
        self.cache.set(query, results)

        return results