메모리 기초

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