Agent Memory 시스템

from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from langchain_databricks import ChatDatabricks

# 모델 설정
model = ChatDatabricks(endpoint="databricks-claude-sonnet-4")

# 체크포인트로 대화 상태 자동 저장
checkpointer = MemorySaver()
agent = create_react_agent(model, tools=[], checkpointer=checkpointer)

# thread_id로 대화 세션 구분 — 같은 thread_id면 이전 대화를 기억
config = {"configurable": {"thread_id": "user-123-session-1"}}

# 첫 번째 대화
result1 = agent.invoke(
    {"messages": [("user", "저는 데이터 엔지니어 김철수입니다.")]},
    config
)

# 두 번째 대화 — 같은 thread_id이므로 이전 대화를 기억
result2 = agent.invoke(
    {"messages": [("user", "제 이름이 뭐였죠?")]},
    config
)
# → "김철수님이시죠" (이전 대화를 기억)

from langchain_core.messages import SystemMessage

def summarize_and_compress(messages, model, max_tokens=50000):
    """오래된 메시지를 요약하여 Context Window 공간 확보"""
    if estimate_tokens(messages) < max_tokens:
        return messages

    # 오래된 메시지들을 요약
    old_messages = messages[:-10]  # 최근 10개는 유지
    recent_messages = messages[-10:]

    summary = model.invoke([
        SystemMessage(content="아래 대화를 핵심 정보만 간결하게 요약하세요."),
        *old_messages
    ])

    # 요약본 + 최근 메시지로 교체
    return [
        SystemMessage(content=f"[이전 대화 요약]\n{summary.content}"),
        *recent_messages
    ]

def sliding_window(messages, window_size=20):
    """최근 N개 메시지만 유지"""
    system_msgs = [m for m in messages if isinstance(m, SystemMessage)]
    conversation = [m for m in messages if not isinstance(m, SystemMessage)]
    return system_msgs + conversation[-window_size:]

def selective_retention(messages, model, max_tokens=50000):
    """각 메시지의 중요도를 평가하여 선별 유지"""
    scored = []
    for msg in messages:
        # 숫자, 이름, 날짜 등 핵심 정보 포함 여부로 점수 부여
        score = calculate_importance(msg)
        scored.append((score, msg))

    scored.sort(key=lambda x: x[0], reverse=True)

    retained = []
    total_tokens = 0
    for score, msg in scored:
        tokens = estimate_tokens([msg])
        if total_tokens + tokens <= max_tokens:
            retained.append(msg)
            total_tokens += tokens

    # 시간순 재정렬
    return sort_by_timestamp(retained)

from databricks.vector_search.client import VectorSearchClient
from datetime import datetime

vsc = VectorSearchClient()
index = vsc.get_index(
    endpoint_name="memory_vs_endpoint",
    index_name="catalog.schema.memory_index"
)

def store_memory(user_id: str, content: str, metadata: dict = None):
    """대화 내용을 임베딩하여 Vector Search에 저장"""
    timestamp = datetime.now().isoformat()
    record = {
        "id": f"{user_id}_{timestamp}",
        "text": content,
        "user_id": user_id,
        "timestamp": timestamp,
        "memory_type": "conversation",
    }
    if metadata:
        record.update(metadata)
    index.upsert([record])

def recall_memory(user_id: str, query: str, top_k: int = 5):
    """관련 기억을 의미 검색으로 가져오기"""
    results = index.similarity_search(
        query_text=query,
        columns=["text", "timestamp", "memory_type"],
        filters={"user_id": user_id},
        num_results=top_k
    )
    return results.get("result", {}).get("data_array", [])

# 사용 예시
store_memory("user-456", "고객이 환불 정책에 대해 문의함. 30일 이내 전액 환불 안내 완료.")
memories = recall_memory("user-456", "환불 관련 이전 문의")

import psycopg2
from datetime import datetime

# Lakebase 연결 (Databricks 서버리스 PostgreSQL)
conn = psycopg2.connect(
    host="<lakebase-endpoint>.cloud.databricks.com",
    port=5432,
    dbname="agent_memory_db",
    user="token",
    password="<databricks-token>"
)
cursor = conn.cursor()

def store_user_preference(user_id: str, key: str, value: str):
    """사용자 프로필 메모리 저장 (UPSERT)"""
    cursor.execute("""
        INSERT INTO agent_memory (user_id, key, value, updated_at)
        VALUES (%s, %s, %s, NOW())
        ON CONFLICT (user_id, key)
        DO UPDATE SET value = %s, updated_at = NOW()
    """, (user_id, key, value, value))
    conn.commit()

def get_user_preferences(user_id: str) -> list:
    """사용자의 모든 저장된 선호/프로필 조회"""
    cursor.execute(
        "SELECT key, value FROM agent_memory WHERE user_id = %s ORDER BY updated_at DESC",
        (user_id,)
    )
    return cursor.fetchall()

# 사용 예시
store_user_preference("user-456", "preferred_language", "한국어")
store_user_preference("user-456", "department", "데이터 엔지니어링팀")
store_user_preference("user-456", "timezone", "Asia/Seoul")

prefs = get_user_preferences("user-456")
# → [("timezone", "Asia/Seoul"), ("department", "데이터 엔지니어링팀"), ...]

┌─────────────────────────────────────────────────────┐
│                 Agent System Prompt                  │
│                                                     │
│  [Profile Memory]    ← Lakebase (구조화 데이터)       │
│  이름: 김철수, 부서: DE팀, 언어: 한국어                │
│                                                     │
│  [Recent Context]    ← Short-term (현재 대화)         │
│  최근 3개 메시지                                      │
│                                                     │
│  [Relevant History]  ← Vector Search (의미 검색)      │
│  과거 유사한 대화 3건                                  │
└─────────────────────────────────────────────────────┘

import mlflow
from mlflow.client import MlflowClient

client = MlflowClient()

def search_successful_episodes(query: str, top_k: int = 3) -> list:
    """과거 성공한 에피소드를 검색하여 참고"""
    # MLflow에서 성공한 실행 이력 조회
    runs = client.search_runs(
        experiment_ids=["<experiment_id>"],
        filter_string='attributes.status = "FINISHED" AND tags.outcome = "success"',
        order_by=["metrics.user_satisfaction DESC"],
        max_results=top_k
    )

    episodes = []
    for run in runs:
        episodes.append({
            "task": run.data.tags.get("task_description", ""),
            "approach": run.data.tags.get("approach", ""),
            "tools_used": run.data.tags.get("tools_used", ""),
            "outcome": run.data.tags.get("outcome", ""),
            "duration_sec": run.data.metrics.get("duration_sec", 0),
        })
    return episodes

def record_episode(task: str, approach: str, tools: list, outcome: str):
    """현재 에피소드를 기록하여 미래 참조용으로 저장"""
    with mlflow.start_run(tags={
        "task_description": task,
        "approach": approach,
        "tools_used": ",".join(tools),
        "outcome": outcome,
    }) as run:
        mlflow.log_metric("user_satisfaction", 1.0 if outcome == "success" else 0.0)

# 사용 예시: 유사 에피소드 검색 후 프롬프트에 주입
past_episodes = search_successful_episodes("ETL 파이프라인 오류 해결")
episode_context = "\n".join([
    f"- 과거 사례: {ep['task']} → {ep['approach']} (결과: {ep['outcome']})"
    for ep in past_episodes
])
# → System Prompt에 [Past Experiences] 섹션으로 삽입

[사용자 요청] → [유사 과거 에피소드 검색]
                        ↓
               성공 에피소드 발견?
              ┌── Yes ──┐── No ──┐
              ↓                   ↓
      성공 경로 우선 시도     기본 추론으로 진행
              ↓                   ↓
         결과 기록 ←──────── 결과 기록
              ↓
      [새 에피소드로 저장]

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, END
import operator

class AgentState(TypedDict):
    """Agent의 Working Memory를 State로 관리"""
    messages: Annotated[list, operator.add]  # 대화 히스토리
    plan: list[str]                           # 현재 실행 계획
    current_step: int                         # 진행 중인 단계
    intermediate_results: dict                # 중간 결과 저장
    scratchpad: str                           # Agent의 메모장

def planning_node(state: AgentState) -> dict:
    """작업 계획을 수립하고 Working Memory에 저장"""
    plan = generate_plan(state["messages"][-1])
    return {
        "plan": plan,
        "current_step": 0,
        "scratchpad": f"총 {len(plan)}단계 계획 수립 완료"
    }

def execution_node(state: AgentState) -> dict:
    """현재 단계를 실행하고 중간 결과를 Working Memory에 기록"""
    step = state["plan"][state["current_step"]]
    result = execute_step(step, state["intermediate_results"])

    new_results = {** state["intermediate_results"], f"step_{state['current_step']}": result}
    return {
        "current_step": state["current_step"] + 1,
        "intermediate_results": new_results,
        "scratchpad": state["scratchpad"] + f"\n단계 {state['current_step']} 완료: {result}"
    }

# 그래프 구성
graph = StateGraph(AgentState)
graph.add_node("plan", planning_node)
graph.add_node("execute", execution_node)

def build_scratchpad_prompt(state: AgentState) -> str:
    """Working Memory의 scratchpad를 프롬프트에 반영"""
    scratchpad = state.get("scratchpad", "")
    intermediate = state.get("intermediate_results", {})

    prompt_section = "[Working Memory - 현재 진행 상황]\n"
    if scratchpad:
        prompt_section += f"메모: {scratchpad}\n"
    if intermediate:
        prompt_section += "중간 결과:\n"
        for key, val in intermediate.items():
            prompt_section += f"  - {key}: {val}\n"

    return prompt_section

def consolidate_memories(user_id: str, days_old: int = 30):
    """30일 이상 된 세부 기억을 요약하여 압축"""
    old_memories = get_memories_older_than(user_id, days_old)

    if not old_memories:
        return

    # LLM으로 오래된 기억들을 요약
    summary = model.invoke(f"""
    다음은 사용자 {user_id}의 과거 대화 기록입니다.
    핵심 정보만 추출하여 간결하게 요약하세요:
    {format_memories(old_memories)}
    """)

    # 요약본 저장 + 원본 삭제
    store_memory(user_id, summary.content, {"memory_type": "consolidated"})
    delete_memories(old_memories)

def forget_user_data(user_id: str):
    """GDPR 삭제 요청 대응: 특정 사용자의 모든 메모리 삭제"""
    # Vector Search에서 삭제
    index.delete(filter={"user_id": user_id})

    # Lakebase에서 삭제
    cursor.execute("DELETE FROM agent_memory WHERE user_id = %s", (user_id,))
    conn.commit()

    # MLflow 에피소드에서 태그 익명화
    anonymize_episodes(user_id)

def recall_with_ranking(user_id: str, query: str, top_k: int = 5):
    """복합 랭킹으로 최적의 메모리 검색"""
    # 1단계: Vector similarity로 후보 검색 (넉넉하게)
    candidates = recall_memory(user_id, query, top_k=top_k * 3)

    # 2단계: 복합 점수 계산
    scored = []
    for mem in candidates:
        similarity_score = mem["score"]  # 의미적 유사도
        recency_score = calculate_recency_weight(mem["timestamp"])  # 최신성
        importance_score = mem.get("importance", 0.5)  # 메타데이터 중요도

        # 가중 합산
        final_score = (
            0.5 * similarity_score +
            0.3 * recency_score +
            0.2 * importance_score
        )
        scored.append((final_score, mem))

    # 3단계: 최종 순위로 반환
    scored.sort(key=lambda x: x[0], reverse=True)
    return [mem for _, mem in scored[:top_k]]

from databricks.agents import ChatAgent

class MyAgent(ChatAgent):
    def predict(self, messages, context=None):
        # messages에 이전 대화 히스토리가 자동 포함됨
        # → 별도 메모리 구현 불필요 (세션 내 한정)
        response = self.model.invoke(messages)
        return response

-- Lakebase에 사용자 프로필 테이블 생성
CREATE TABLE agent_memory (
    user_id     VARCHAR(255) NOT NULL,
    key         VARCHAR(255) NOT NULL,
    value       TEXT,
    created_at  TIMESTAMP DEFAULT NOW(),
    updated_at  TIMESTAMP DEFAULT NOW(),
    expires_at  TIMESTAMP,  -- TTL 지원
    PRIMARY KEY (user_id, key)
);

-- 인덱스 추가
CREATE INDEX idx_memory_user ON agent_memory(user_id);
CREATE INDEX idx_memory_expires ON agent_memory(expires_at);

-- Delta Table로 메모리 저장소 생성
CREATE TABLE catalog.schema.conversation_memory (
    id          STRING,
    user_id     STRING,
    content     STRING,
    timestamp   TIMESTAMP,
    memory_type STRING,
    importance  DOUBLE DEFAULT 0.5
);

-- Vector Search Index 생성
CREATE VECTOR SEARCH INDEX memory_index
ON catalog.schema.conversation_memory (content)
USING MODEL `databricks-gte-large-en`
WITH (endpoint_name = "memory_vs_endpoint");

def build_memory_enhanced_prompt(user_id: str, current_query: str) -> str:
    """모든 메모리 소스를 결합하여 풍부한 프롬프트 생성"""

    # 1. User Profile Memory (Lakebase)
    profile = get_user_preferences(user_id)
    profile_section = "[사용자 프로필]\n"
    for key, value in profile:
        profile_section += f"- {key}: {value}\n"

    # 2. Relevant History (Vector Search)
    relevant_memories = recall_with_ranking(user_id, current_query, top_k=3)
    history_section = "[관련 과거 대화]\n"
    for mem in relevant_memories:
        history_section += f"- ({mem['timestamp']}): {mem['text']}\n"

    # 3. Episodic Memory (MLflow)
    episodes = search_successful_episodes(current_query, top_k=2)
    episode_section = "[참고할 과거 경험]\n"
    for ep in episodes:
        episode_section += f"- {ep['task']}: {ep['approach']} → {ep['outcome']}\n"

    # 통합 System Prompt
    system_prompt = f"""당신은 전문 업무 지원 Agent입니다.

{profile_section}
{history_section}
{episode_section}

위 메모리를 참고하여 사용자에게 개인화되고 맥락에 맞는 답변을 제공하세요.
새롭게 파악된 중요 정보는 메모리에 저장하세요."""

    return system_prompt