에이전트 메모리 컴팩션, 직접 구현하기 — 멀티턴 루프부터 캐시·장기 메모리까지

import anthropic
 
client = anthropic.Anthropic()  # ANTHROPIC_API_KEY 환경변수 사용
 
MODEL = "claude-sonnet-4-6"
SYSTEM = "너는 데이터 플랫폼 운영을 돕는 시니어 엔지니어다. 한국어로 간결하게 답한다."
 
# 이 리스트가 곧 '단기 기억'이다. 그 이상도 이하도 아니다.
messages: list[dict] = []
 
def chat(user_text: str) -> str:
    messages.append({"role": "user", "content": user_text})
 
    resp = client.messages.create(
        model=MODEL,
        max_tokens=1024,
        system=SYSTEM,
        messages=messages,          # ← 매 턴 '전체'를 다시 보낸다
    )
 
    assistant_text = "".join(
        block.text for block in resp.content if block.type == "text"
    )
    messages.append({"role": "assistant", "content": assistant_text})
    return assistant_text

def count_tokens(messages: list[dict], system: str, tools: list | None = None) -> int:
    res = client.messages.count_tokens(
        model=MODEL,
        system=system,
        messages=messages,
        tools=tools or [],
    )
    return res.input_tokens

# create() 응답에는 항상 usage가 들어 있다 — 추가 호출 없이 점유량을 안다.
resp = client.messages.create(...)
used = resp.usage.input_tokens + resp.usage.output_tokens

CONTEXT_WINDOW = 200_000   # 사용하는 모델의 컨텍스트 창
COMPACT_AT = 0.75          # 75% 차면 컴팩션
KEEP_RECENT_TOKENS = 20_000  # 최근 이만큼은 원문 보존

def split_old_recent(messages: list[dict], keep_recent_tokens: int):
    """뒤에서부터 누적하여 최근 부분을 확보하고, 나머지를 오래된 부분으로 반환."""
    recent: list[dict] = []
    acc = 0
    # 뒤(최근)에서 앞(과거)으로 훑는다
    for msg in reversed(messages):
        t = count_tokens([msg], system="")  # 메시지 단위 근사 측정
        if acc + t > keep_recent_tokens and recent:
            break
        recent.append(msg)
        acc += t
    recent.reverse()
    split_idx = len(messages) - len(recent)
    old = messages[:split_idx]
    return old, recent

def fix_tool_boundary(old: list[dict], recent: list[dict]):
    """recent의 첫 메시지가 tool_result로 시작하면, 짝이 되는 tool_use가
    old의 마지막에 있다는 뜻이다. 경계를 한 턴 당겨 짝을 보존한다."""
    while recent and _starts_with_tool_result(recent[0]):
        old = old + [recent.pop(0)]  # 어차피 이 한 줄도 곧 요약된다
    return old, recent
 
def _starts_with_tool_result(msg: dict) -> bool:
    content = msg.get("content")
    if isinstance(content, list) and content:
        return content[0].get("type") == "tool_result"
    return False

COMPACT_PROMPT = """\
아래는 한 에이전트 세션의 지난 대화 기록이다. 이어서 작업을 계속할 수 있도록
핵심만 추려 '인수인계 메모'를 작성하라. 반드시 아래 항목 구조를 그대로 따르고,
해당 없는 항목은 "없음"이라고 적어라. 추측·창작 금지, 기록에 있는 사실만.
 
## 사용자의 목표
## 확정된 결정·제약·선호
## 현재 진행 중인 작업(가장 중요)
## 알아낸 핵심 사실(이름/경로/버전/수치/식별자)
## 시도했다 폐기한 막다른 길(반복 방지용)
## 다음에 할 일
"""
 
def summarize_old(old: list[dict]) -> str:
    res = client.messages.create(
        model=MODEL,
        max_tokens=2048,
        system="너는 대화 기록을 인수인계 메모로 압축하는 요약기다.",
        messages=old + [{"role": "user", "content": COMPACT_PROMPT}],
    )
    return "".join(b.text for b in res.content if b.type == "text")

def validate_summary(text: str) -> bool:
    # 양식이 유지됐고, 진행 중 작업 칸이 통째로 비지 않았는지 최소 확인
    must_have = ["## 현재 진행 중인 작업", "## 다음에 할 일"]
    return all(h in text for h in must_have)

def compact(messages: list[dict]) -> list[dict]:
    old, recent = split_old_recent(messages, KEEP_RECENT_TOKENS)
    if not old:
        return messages  # 압축할 과거가 없음
 
    old, recent = fix_tool_boundary(old, recent)
    summary = summarize_old(old)
    if not validate_summary(summary):
        summary = summarize_old(old)  # 1회 재시도
 
    marker = (
        "[이전 대화 요약 — 원문은 컨텍스트 절약을 위해 압축됨]\n\n" + summary
    )
    return [{"role": "user", "content": marker}] + recent

TOOL_RESULT_MAX_CHARS = 4000
 
def shrink_tool_result(name: str, args: dict, raw: str) -> str:
    if len(raw) <= TOOL_RESULT_MAX_CHARS:
        return raw
    head = raw[:1500]
    tail = raw[-1500:]
    pointer = f"{name}({args})"  # 동일 인자로 다시 부르면 전체를 재취득 가능
    return (
        f"{head}\n"
        f"... [중략: 전체 {len(raw)}자 중 일부만 보존. "
        f"전체가 다시 필요하면 `{pointer}`를 재실행할 것] ...\n"
        f"{tail}"
    )

def run_tool_and_record(messages, tool_use_block, registry):
    name = tool_use_block.name
    args = tool_use_block.input
    raw = registry[name](**args)              # 실제 도구 실행
    shrunk = shrink_tool_result(name, args, raw)
    messages.append({
        "role": "user",
        "content": [{
            "type": "tool_result",
            "tool_use_id": tool_use_block.id,
            "content": shrunk,                # ← 축약본을 기억에 넣는다
        }],
    })

resp = client.messages.create(
    model=MODEL,
    max_tokens=1024,
    system=[
        {"type": "text", "text": SYSTEM},
        # 시스템 프롬프트는 거의 안 변하므로 캐시하기 좋다
        {"type": "text", "text": LONG_STABLE_GUIDE,
         "cache_control": {"type": "ephemeral"}},
    ],
    messages=messages,
)
# 캐시 효과는 usage로 확인한다
print(resp.usage.cache_creation_input_tokens,  # 이번에 새로 캐시한 양
      resp.usage.cache_read_input_tokens)       # 캐시에서 재사용한 양

def with_cache_breakpoint(messages: list[dict]) -> list[dict]:
    """맨 앞 요약 메시지에 캐시 경계를 부여한다(컴팩션 직후 호출)."""
    if not messages:
        return messages
    first = messages[0]
    # content를 블록 리스트로 정규화한 뒤 마지막 블록에 cache_control 부여
    blocks = (first["content"] if isinstance(first["content"], list)
              else [{"type": "text", "text": first["content"]}])
    blocks[-1] = {**blocks[-1], "cache_control": {"type": "ephemeral"}}
    return [{**first, "content": blocks}] + messages[1:]

import json, pathlib
 
MEM_DIR = pathlib.Path("./agent_memory")
MEM_DIR.mkdir(exist_ok=True)
 
def remember(key: str, fact: str):
    (MEM_DIR / f"{key}.json").write_text(
        json.dumps({"key": key, "fact": fact}, ensure_ascii=False))
 
def load_long_term() -> str:
    facts = []
    for f in sorted(MEM_DIR.glob("*.json")):
        d = json.loads(f.read_text())
        facts.append(f"- {d['fact']}")
    return "\n".join(facts)
 
# 세션 시작 시: 장기 기억을 시스템 프롬프트에 주입
def build_system() -> str:
    longterm = load_long_term()
    if not longterm:
        return SYSTEM
    return SYSTEM + "\n\n[이전 세션에서 알게 된 사실]\n" + longterm

def recall_relevant(query: str, k: int = 5) -> str:
    hits = vector_store.search(embed(query), top_k=k)  # 임베딩 유사도 검색
    return "\n".join(f"- {h.text}" for h in hits)
 
def chat_with_recall(user_text: str):
    # 이번 턴에만 필요한 기억을 끌어와 일시적으로 끼운다
    relevant = recall_relevant(user_text)
    augmented = user_text
    if relevant:
        augmented = f"[참고 기억]\n{relevant}\n\n[질문]\n{user_text}"
    messages.append({"role": "user", "content": augmented})
    # ... 이하 동일

class CompactingAgent:
    def __init__(self, tools: dict):
        self.tools = tools
        self.messages: list[dict] = []
 
    def step(self, user_text: str) -> str:
        self.messages.append({"role": "user", "content": user_text})
 
        while True:
            resp = client.messages.create(
                model=MODEL,
                max_tokens=2048,
                system=build_system(),                 # 장기 기억 주입
                messages=self.messages,
                tools=[t.schema for t in self.tools.values()],
            )
            self.messages.append({"role": "assistant", "content": resp.content})
 
            if resp.stop_reason == "tool_use":
                for block in resp.content:
                    if block.type == "tool_use":
                        run_tool_and_record(self.messages, block, self.tools)
                self._maybe_compact()                  # 도구 후에도 점검
                continue                               # 도구 결과로 한 번 더
 
            self._maybe_compact()                      # 턴 종료 후 점검
            return "".join(b.text for b in resp.content if b.type == "text")
 
    def _maybe_compact(self):
        used = count_tokens(self.messages, build_system(),
                            [t.schema for t in self.tools.values()])
        if used / CONTEXT_WINDOW > COMPACT_AT:
            self.messages = compact(self.messages)
            self.messages = with_cache_breakpoint(self.messages)