from langchain.memory import ConversationSummaryBufferMemory
from langchain_openai import ChatOpenAI
from langchain.chains import ConversationChain

# 初始化 LLM
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# SummaryBufferMemory：近 2000 Token 保留原文，超出则压缩为摘要
memory = ConversationSummaryBufferMemory(
    llm=llm,
    max_token_limit=2000,        # 超过此阈值触发摘要压缩
    memory_key="chat_history",
    return_messages=True,
    human_prefix="User",
    ai_prefix="Assistant"
)

# 构建对话链
chain = ConversationChain(
    llm=llm,
    memory=memory,
    verbose=True
)

# 模拟多轮对话
response1 = chain.predict(input="我是小明，我在开发一个电商平台")
response2 = chain.predict(input="我主要负责后端 Java Spring Boot 模块")
response3 = chain.predict(input="现在帮我设计用户订单状态机")

# 查看当前记忆内容
print("== 当前记忆摘要 ==")
print(memory.moving_summary_buffer)   # 已压缩的历史摘要
print("== 近期原文记忆 ==")
for msg in memory.chat_memory.messages:
    print(f"{msg.type}: {msg.content[:100]}")

import json
import redis
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationEntityMemory
from langchain.memory.entity import RedisEntityStore

# Redis 实体存储后端（生产推荐，支持 TTL 和持久化）
redis_client = redis.Redis(host="localhost", port=6379, db=0)
entity_store = RedisEntityStore(
    redis_client=redis_client,
    session_id="user_alice_session",  # 按用户隔离
    key_prefix="entity",
    ttl=86400 * 30            # 实体信息保留30天
)

llm = ChatOpenAI(model="gpt-4o-mini")

memory = ConversationEntityMemory(
    llm=llm,
    entity_store=entity_store,
    memory_key="chat_history",
    input_key="input",
    k=5,                        # 注入最近5条相关实体
    return_messages=True
)

# 第一次会话：建立用户画像
memory.save_context(
    {"input": "我叫 Alice，是字节跳动的 Java 工程师，负责推荐系统"},
    {"output": "好的 Alice，很高兴认识你！推荐系统有什么我可以帮你的？"}
)

# 查看抽取到的实体
entities = entity_store.mget(["Alice", "字节跳动"])
for k, v in zip(["Alice", "字节跳动"], entities):
    print(f"{k}: {v}")
# 输出示例:
# Alice: Java工程师，就职于字节跳动，负责推荐系统开发
# 字节跳动: Alice的雇主，互联网公司

# === 三天后，新会话 ===
# 重新加载同一 session_id 的 entity_store，实体信息仍然存在
new_memory = ConversationEntityMemory(
    llm=llm,
    entity_store=RedisEntityStore(
        redis_client=redis_client,
        session_id="user_alice_session"  # 相同 session_id
    )
)
# Agent 仍然"记得"Alice是Java工程师，负责推荐系统
result = new_memory.load_memory_variables({"input": "帮我优化一下推荐算法"})

import time
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, Range
from openai import OpenAI
import uuid

openai_client = OpenAI()
qdrant = QdrantClient(url="http://localhost:6333")

# 创建记忆集合
COLLECTION = "agent_memory"
qdrant.recreate_collection(
    collection_name=COLLECTION,
    vectors_config=VectorParams(size=1536, distance=Distance.COSINE)
)

def embed(text: str) -> list[float]:
    """调用 OpenAI text-embedding-3-small"""
    resp = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return resp.data[0].embedding

def save_memory(user_id: str, text: str, importance: float = 1.0):
    """将一条记忆写入向量库"""
    vector = embed(text)
    qdrant.upsert(
        collection_name=COLLECTION,
        points=[PointStruct(
            id=str(uuid.uuid4()),
            vector=vector,
            payload={
                "user_id": user_id,
                "text": text,
                "timestamp": time.time(),     # Unix 时间戳，用于时间衰减
                "importance": importance       # 0~1，重要性评分
            }
        )]
    )

def retrieve_memory(user_id: str, query: str, top_k: int = 5) -> list[dict]:
    """检索最相关记忆，结合语义相似度与时间衰减"""
    vector = embed(query)
    results = qdrant.search(
        collection_name=COLLECTION,
        query_vector=vector,
        query_filter=Filter(
            must=[FieldCondition(
                key="user_id",
                match={"value": user_id}   # 用户级隔离
            )]
        ),
        limit=top_k * 3,               # 多取一些，后续重排序
        with_payload=True
    )

    now = time.time()
    scored = []
    for r in results:
        age_days = (now - r.payload["timestamp"]) / 86400
        # 时间衰减：7天后相关性减半，利用指数衰减
        time_decay = 0.5 ** (age_days / 7)
        final_score = r.score * 0.7 + time_decay * 0.2 + r.payload["importance"] * 0.1
        scored.append({"text": r.payload["text"], "score": final_score})

    scored.sort(key=lambda x: x["score"], reverse=True)
    return scored[:top_k]

# 写入几条历史记忆
save_memory("alice", "Alice偏好使用 Java 开发后端，擅长 Spring Boot 和 Kafka", importance=0.9)
save_memory("alice", "Alice在2026年3月完成了推荐系统 A/B 测试，CTR 提升了12%", importance=0.8)
save_memory("alice", "Alice的团队正在迁移到 K8s，遇到了 Ingress 配置问题", importance=0.7)

# 检索相关记忆
memories = retrieve_memory("alice", "帮我设计一个高并发订单服务")
context = "\n".join([m["text"] for m in memories])
print("注入的记忆上下文:\n", context)

from typing import Annotated, TypedDict, Optional
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
import operator

# ── State Schema：三层记忆统一管理 ──
class AgentState(TypedDict):
    # Layer 1: 当前会话消息列表（LangGraph 自动追加）
    messages: Annotated[list[BaseMessage], operator.add]

    # Layer 2: 实体记忆（结构化 KV）
    entities: dict[str, str]           # {"Alice": "Java工程师...", "项目X": "..."}

    # Layer 3: 语义记忆检索结果（注入当前轮）
    retrieved_memories: list[str]

    # 用户标识（用于记忆隔离）
    user_id: str

# ── 节点1：从向量库检索相关记忆 ──
def retrieve_memories_node(state: AgentState) -> dict:
    latest_msg = state["messages"][-1].content
    memories = retrieve_memory(state["user_id"], latest_msg, top_k=3)
    return {"retrieved_memories": [m["text"] for m in memories]}

# ── 节点2：LLM 推理（注入三层记忆）──
def agent_node(state: AgentState) -> dict:
    # 构建 System Prompt（融合三层记忆）
    entity_ctx = "\n".join([f"{k}: {v}" for k, v in state["entities"].items()])
    semantic_ctx = "\n".join(state["retrieved_memories"])

    system_prompt = f"""你是一个专业的 AI 助手。

【实体记忆（用户背景）】
{entity_ctx}

【相关历史经验】
{semantic_ctx}

请结合以上背景信息回答用户问题。"""

    # 调用 LLM（此处省略具体实现）
    response = llm.invoke([{"role": "system", "content": system_prompt}] + state["messages"])
    return {"messages": [AIMessage(content=response.content)]}

# ── 节点3：更新记忆（对话后写入）──
def update_memory_node(state: AgentState) -> dict:
    # 异步写入语义记忆（最近一轮对话）
    last_human = next(m for m in reversed(state["messages"]) if isinstance(m, HumanMessage))
    last_ai = next(m for m in reversed(state["messages"]) if isinstance(m, AIMessage))
    memory_text = f"用户问：{last_human.content}\n助手答：{last_ai.content[:200]}"
    save_memory(state["user_id"], memory_text)
    return {}

# ── 构建图 ──
graph = StateGraph(AgentState)
graph.add_node("retrieve", retrieve_memories_node)
graph.add_node("agent", agent_node)
graph.add_node("update_memory", update_memory_node)

graph.set_entry_point("retrieve")
graph.add_edge("retrieve", "agent")
graph.add_edge("agent", "update_memory")
graph.add_edge("update_memory", END)

# 持久化 Checkpointing（支持会话恢复）
checkpointer = SqliteSaver.from_conn_string("agent_memory.db")
app = graph.compile(checkpointer=checkpointer)

# 调用 Agent（thread_id 用于识别会话）
result = app.invoke(
    {
        "messages": [HumanMessage(content="帮我设计一个高并发订单服务")],
        "entities": {"Alice": "Java工程师，字节跳动，推荐系统"},
        "retrieved_memories": [],
        "user_id": "alice"
    },
    config={"configurable": {"thread_id": "alice_thread_001"}}
)