Memory 시스템

from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.graph import StateGraph, END, MessagesState
from langchain_anthropic import ChatAnthropic
import os

async def create_persistent_agent():
    # Postgres에 대화 상태 영구 저장
    checkpointer = await AsyncPostgresSaver.from_conn_string(
        os.getenv("DATABASE_URL")
    )
    await checkpointer.setup()  # 테이블 초기화

    llm = ChatAnthropic(model="claude-opus-4-7")

    async def agent_node(state: MessagesState):
        response = await llm.ainvoke(state["messages"])
        return {"messages": [response]}

    builder = StateGraph(MessagesState)
    builder.add_node("agent", agent_node)
    builder.set_entry_point("agent")
    builder.add_edge("agent", END)

    return builder.compile(checkpointer=checkpointer)

# 사용: thread_id로 대화 세션 구분
async def chat(agent, user_id: str, message: str) -> str:
    config = {"configurable": {"thread_id": user_id}}

    result = await agent.ainvoke(
        {"messages": [{"role": "user", "content": message}]},
        config=config
    )
    return result["messages"][-1].content

# 사용자별 대화 이력 조회
async def get_conversation_history(agent, user_id: str) -> list:
    config = {"configurable": {"thread_id": user_id}}
    state = await agent.aget_state(config)
    return state.values.get("messages", [])

from langchain_chroma import Chroma
from langchain_anthropic import ChatAnthropic
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
from datetime import datetime

class SemanticMemory:
    def __init__(self, persist_dir: str = "./memory_store"):
        self.vectorstore = Chroma(
            embedding_function=OpenAIEmbeddings(),
            persist_directory=persist_dir,
            collection_name="agent_memory"
        )
        self.llm = ChatAnthropic(model="claude-haiku-4-5-20251001")

    async def remember(self, content: str, user_id: str, metadata: dict = {}) -> str:
        """대화에서 기억할 내용 추출 + 저장"""
        # LLM으로 기억할 핵심 정보 추출
        extraction = await self.llm.ainvoke([{
            "role": "user",
            "content": f"다음 대화에서 사용자에 대해 기억할 중요한 사실을 한 문장으로 추출하세요. 없으면 'SKIP':\n{content}"
        }])
        fact = extraction.content.strip()
        if fact == "SKIP":
            return

        doc = Document(
            page_content=fact,
            metadata={
                "user_id": user_id,
                "timestamp": datetime.now().isoformat(),
                **metadata
            }
        )
        self.vectorstore.add_documents([doc])
        return fact

    async def recall(self, query: str, user_id: str, k: int = 3) -> list[str]:
        """유사한 기억 검색"""
        docs = self.vectorstore.similarity_search(
            query,
            k=k,
            filter={"user_id": user_id}
        )
        return [d.page_content for d in docs]

    async def build_memory_context(self, query: str, user_id: str) -> str:
        """검색된 기억을 시스템 프롬프트용 텍스트로 변환"""
        memories = await self.recall(query, user_id)
        if not memories:
            return ""
        memory_list = "\n".join(f"- {m}" for m in memories)
        return f"\n[이 사용자에 대해 알고 있는 것]\n{memory_list}\n"

# Agent에서 사용
memory = SemanticMemory()

async def memory_agent_node(state):
    user_id = state["user_id"]
    query = state["messages"][-1].content

    # 관련 기억 검색
    mem_context = await memory.build_memory_context(query, user_id)

    llm = ChatAnthropic(model="claude-opus-4-7")
    response = await llm.ainvoke(
        state["messages"],
        system=f"당신은 개인화된 AI 어시스턴트입니다.{mem_context}"
    )

    # 새 기억 저장
    await memory.remember(query, user_id)
    return {"messages": [response]}

from langchain_core.messages import SystemMessage, HumanMessage, AIMessage

MAX_MESSAGES = 20     # 유지할 최근 메시지 수
SUMMARY_KEEP = 6      # 압축 후 유지할 최신 메시지 수

async def compress_messages(
    messages: list,
    llm: ChatAnthropic
) -> list:
    """메시지가 MAX_MESSAGES를 초과하면 요약 압축"""
    if len(messages) <= MAX_MESSAGES:
        return messages

    # 오래된 메시지를 요약
    old_messages = messages[:-SUMMARY_KEEP]
    recent_messages = messages[-SUMMARY_KEEP:]

    summary_response = await llm.ainvoke([
        SystemMessage(content="다음 대화를 핵심 정보 위주로 간결하게 요약하세요."),
        *old_messages
    ])
    summary_text = summary_response.content

    # 요약 + 최신 메시지 결합
    compressed = [
        SystemMessage(content=f"[이전 대화 요약]\n{summary_text}"),
        *recent_messages
    ]
    print(f"메시지 압축: {len(messages)} → {len(compressed)}")
    return compressed

# LangGraph 노드에서 사용
async def agent_with_compression(state):
    messages = state["messages"]
    llm = ChatAnthropic(model="claude-haiku-4-5-20251001")

    # 압축 후 LLM 호출
    compressed = await compress_messages(messages, llm)
    response = await ChatAnthropic(model="claude-opus-4-7").ainvoke(compressed)
    return {"messages": messages + [response]}  # 원본 상태에는 전체 보존