메모리와 상태 관리 — 에이전트가 기억하는 방법

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

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

def summarize_conversation(messages: list) -> str:
    """대화 기록을 요약해 컨텍스트를 줄입니다."""
    conversation_text = "\n".join([
        f"{m.type}: {m.content}" for m in messages
    ])
    
    summary = llm.invoke([
        SystemMessage("다음 대화를 핵심만 3문장으로 요약해줘."),
        HumanMessage(conversation_text)
    ])
    return summary.content

def trim_messages(messages: list, max_messages: int = 20) -> list:
    """메시지가 너무 많으면 오래된 것을 요약으로 대체합니다."""
    if len(messages) <= max_messages:
        return messages
    
    # 오래된 메시지 요약
    old_messages = messages[:-max_messages]
    recent_messages = messages[-max_messages:]
    
    summary = summarize_conversation(old_messages)
    
    return [
        SystemMessage(f"이전 대화 요약: {summary}"),
        *recent_messages
    ]

from typing import Annotated, Optional
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    user_name: Optional[str]          # 사용자 이름
    task_context: Optional[str]       # 현재 작업 컨텍스트
    completed_tasks: list             # 완료된 작업 목록

from langchain_core.tools import tool

@tool
def remember_user_info(key: str, value: str) -> str:
    """사용자의 중요 정보를 기억합니다.
    
    Args:
        key: 정보 종류 (예: 'name', 'preference', 'goal')
        value: 저장할 값
    """
    # 실제 구현에서는 DB 저장
    return f"✅ 기억 완료: {key} = {value}"

@tool
def recall_user_info(key: str) -> str:
    """이전에 기억한 사용자 정보를 불러옵니다."""
    # 실제 구현에서는 DB 조회
    return f"{key}에 대한 정보: [저장된 값]"

import chromadb
from chromadb.utils import embedding_functions
from datetime import datetime

chroma_client = chromadb.PersistentClient(path="./agent_memory")
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key=os.getenv("OPENAI_API_KEY"),
    model_name="text-embedding-3-small"
)

memory_collection = chroma_client.get_or_create_collection(
    name="agent_long_term_memory",
    embedding_function=openai_ef
)

def store_memory(content: str, memory_type: str, user_id: str):
    """중요한 정보를 장기 메모리에 저장합니다."""
    memory_id = f"{user_id}_{datetime.now().timestamp()}"
    
    memory_collection.add(
        ids=[memory_id],
        documents=[content],
        metadatas=[{
            "type": memory_type,   # "fact", "preference", "event"
            "user_id": user_id,
            "timestamp": datetime.now().isoformat()
        }]
    )

def retrieve_memory(query: str, user_id: str, top_k: int = 5) -> list[str]:
    """관련된 장기 기억을 검색합니다."""
    results = memory_collection.query(
        query_texts=[query],
        n_results=top_k,
        where={"user_id": user_id}
    )
    return results["documents"][0] if results["documents"] else []

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode, tools_condition

def agent_with_memory(state: AgentState):
    messages = state["messages"]
    user_id = state.get("user_id", "default")
    
    # 장기 메모리에서 관련 정보 검색
    last_user_message = messages[-1].content if messages else ""
    memories = retrieve_memory(last_user_message, user_id)
    
    # 메모리를 시스템 메시지에 주입
    memory_context = ""
    if memories:
        memory_context = f"\n\n[기억하고 있는 정보]\n" + "\n".join(memories)
    
    system_message = SystemMessage(
        f"당신은 사용자를 도와주는 AI 어시스턴트입니다.{memory_context}"
    )
    
    full_messages = [system_message] + messages
    response = llm_with_tools.invoke(full_messages)
    
    return {"messages": [response]}

from datetime import datetime, timedelta

def forget_old_memories(user_id: str, days_threshold: int = 90):
    """오래되고 중요하지 않은 메모리를 삭제합니다."""
    cutoff = (datetime.now() - timedelta(days=days_threshold)).isoformat()
    
    # 오래된 메모리 조회
    old_memories = memory_collection.get(
        where={
            "user_id": user_id,
            "type": "temporary",  # 임시 메모리만
            "timestamp": {"$lt": cutoff}
        }
    )
    
    if old_memories["ids"]:
        memory_collection.delete(ids=old_memories["ids"])
        return f"{len(old_memories['ids'])}개 오래된 메모리 삭제"
    
    return "삭제할 메모리 없음"

def should_summarize(messages: list, threshold: int = 15) -> bool:
    return len(messages) > threshold

def agent_node_with_summary(state: AgentState):
    messages = state["messages"]
    
    # 메시지가 너무 많으면 요약
    if should_summarize(messages):
        messages = trim_messages(messages)
    
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}