LLM Agent Memory System Design: From Short-Term Memory to Long-Term Memory

from collections import deque
from openai import OpenAI

class ShortTermMemory:
    def __init__(self, max_messages=20):
        self.messages = deque(maxlen=max_messages)
        self.system_prompt = "당신은 도움이 되는 AI 어시스턴트입니다."

    def add(self, role, content):
        self.messages.append({"role": role, "content": content})

    def get_context(self):
        return [{"role": "system", "content": self.system_prompt}] + list(self.messages)

client = OpenAI()
memory = ShortTermMemory(max_messages=10)

def chat(user_input):
    memory.add("user", user_input)
    response = client.chat.completions.create(
        model="gpt-4o", messages=memory.get_context()
    )
    reply = response.choices[0].message.content
    memory.add("assistant", reply)
    return reply

import anthropic

client = anthropic.Anthropic()

class SummaryMemory:
    def __init__(self, window_size=10, summary_threshold=8):
        self.recent_messages = []
        self.summary = ""
        self.window_size = window_size
        self.summary_threshold = summary_threshold

    def _summarize(self):
        to_summarize = self.recent_messages[:self.summary_threshold]
        existing = f"기존 요약: {self.summary}

" if self.summary else ""
        conversation = "
".join(f"{m['role']}: {m['content']}" for m in to_summarize)

        response = client.messages.create(
            model="claude-haiku-4-5-20251001",
            max_tokens=300,
            messages=[{
                "role": "user",
                "content": f"{existing}다음 대화를 3-4문장으로 핵심만 요약하세요:

{conversation}"
            }]
        )
        return response.content[0].text

    def add(self, role, content):
        self.recent_messages.append({"role": role, "content": content})
        if len(self.recent_messages) > self.window_size:
            self.summary = self._summarize()
            self.recent_messages = self.recent_messages[self.summary_threshold:]

    def get_context(self):
        messages = []
        if self.summary:
            messages += [
                {"role": "user", "content": f"[이전 대화 요약]
{self.summary}"},
                {"role": "assistant", "content": "이전 대화 내용을 이해했습니다."}
            ]
        return messages + self.recent_messages

from openai import OpenAI
import json
from datetime import datetime

client = OpenAI()

class LongTermMemory:
    def __init__(self, user_id, supabase_client):
        self.user_id = user_id
        self.sb = supabase_client

    def _embed(self, text):
        response = client.embeddings.create(
            model="text-embedding-3-small", input=text
        )
        return response.data[0].embedding

    def save(self, content, memory_type="fact"):
        embedding = self._embed(content)
        self.sb.table("agent_memories").insert({
            "user_id": self.user_id,
            "content": content,
            "memory_type": memory_type,
            "embedding": embedding,
            "created_at": datetime.utcnow().isoformat()
        }).execute()

    def recall(self, query, top_k=3):
        query_embedding = self._embed(query)
        result = self.sb.rpc("match_memories", {
            "query_embedding": query_embedding,
            "user_id": self.user_id,
            "match_count": top_k
        }).execute()
        return [r["content"] for r in (result.data or [])]

class AgentWithMemory:
    def __init__(self, user_id):
        self.short_term = SummaryMemory(window_size=10)
        self.long_term = LongTermMemory(user_id=user_id)
        self.client = OpenAI()

    def chat(self, user_input):
        # 1. 장기 메모리에서 관련 기억 회상
        relevant_memories = self.long_term.recall(user_input)

        system_content = "당신은 도움이 되는 AI 어시스턴트입니다."
        if relevant_memories:
            system_content += "

[관련 기억]
" + "
".join(f"- {m}" for m in relevant_memories)

        # 2. 단기 + 장기 컨텍스트 조합
        messages = [{"role": "system", "content": system_content}]
        messages.extend(self.short_term.get_context())
        messages.append({"role": "user", "content": user_input})

        response = self.client.chat.completions.create(model="gpt-4o", messages=messages)
        reply = response.choices[0].message.content

        self.short_term.add("user", user_input)
        self.short_term.add("assistant", reply)

        # 3. 중요 정보를 장기 메모리에 저장
        self.long_term.extract_and_save(f"사용자: {user_input}
어시스턴트: {reply}")

        return reply