AutoGPT 메모리 폭주 해결 - SQLite 체크포인트

PRAGMA journal_mode = WAL;
PRAGMA synchronous = NORMAL;
PRAGMA temp_store = MEMORY;
PRAGMA foreign_keys = ON;
PRAGMA busy_timeout = 5000;
PRAGMA wal_autocheckpoint = 0;

PRAGMA wal_checkpoint(TRUNCATE);

CREATE TABLE IF NOT EXISTS runs (
  run_id TEXT PRIMARY KEY,
  created_at INTEGER NOT NULL,
  status TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS steps (
  run_id TEXT NOT NULL,
  step_no INTEGER NOT NULL,
  created_at INTEGER NOT NULL,
  user_input TEXT,
  model_output TEXT,
  tool_name TEXT,
  tool_input TEXT,
  tool_output TEXT,
  summary TEXT,
  tokens_in INTEGER,
  tokens_out INTEGER,
  PRIMARY KEY (run_id, step_no)
);

CREATE TABLE IF NOT EXISTS kv_state (
  run_id TEXT NOT NULL,
  key TEXT NOT NULL,
  value_json TEXT NOT NULL,
  updated_at INTEGER NOT NULL,
  PRIMARY KEY (run_id, key)
);

CREATE INDEX IF NOT EXISTS idx_steps_run_created ON steps(run_id, created_at);

CREATE TABLE IF NOT EXISTS blobs (
  blob_id TEXT PRIMARY KEY,
  run_id TEXT NOT NULL,
  step_no INTEGER NOT NULL,
  kind TEXT NOT NULL,
  content BLOB NOT NULL,
  created_at INTEGER NOT NULL
);

ALTER TABLE steps ADD COLUMN tool_output_blob_id TEXT;

import json
import sqlite3
import time
import uuid
from contextlib import contextmanager

DB_PATH = "agent.db"

@contextmanager
def connect():
    conn = sqlite3.connect(DB_PATH)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")
    conn.execute("PRAGMA foreign_keys=ON")
    conn.execute("PRAGMA busy_timeout=5000")
    conn.execute("PRAGMA wal_autocheckpoint=0")
    try:
        yield conn
        conn.commit()
    finally:
        conn.close()

def init_schema():
    with connect() as conn:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS runs (
              run_id TEXT PRIMARY KEY,
              created_at INTEGER NOT NULL,
              status TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS steps (
              run_id TEXT NOT NULL,
              step_no INTEGER NOT NULL,
              created_at INTEGER NOT NULL,
              user_input TEXT,
              model_output TEXT,
              tool_name TEXT,
              tool_input TEXT,
              tool_output TEXT,
              summary TEXT,
              tokens_in INTEGER,
              tokens_out INTEGER,
              PRIMARY KEY (run_id, step_no)
            );
            CREATE TABLE IF NOT EXISTS kv_state (
              run_id TEXT NOT NULL,
              key TEXT NOT NULL,
              value_json TEXT NOT NULL,
              updated_at INTEGER NOT NULL,
              PRIMARY KEY (run_id, key)
            );
            CREATE INDEX IF NOT EXISTS idx_steps_run_created ON steps(run_id, created_at);
            """
        )

def create_run():
    run_id = str(uuid.uuid4())
    now = int(time.time())
    with connect() as conn:
        conn.execute(
            "INSERT INTO runs(run_id, created_at, status) VALUES(?, ?, ?)",
            (run_id, now, "running"),
        )
    return run_id

def save_step(run_id: str, step_no: int, payload: dict):
    now = int(time.time())
    with connect() as conn:
        conn.execute(
            """
            INSERT OR REPLACE INTO steps(
              run_id, step_no, created_at,
              user_input, model_output,
              tool_name, tool_input, tool_output,
              summary, tokens_in, tokens_out
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                run_id,
                step_no,
                now,
                payload.get("user_input"),
                payload.get("model_output"),
                payload.get("tool_name"),
                payload.get("tool_input"),
                payload.get("tool_output"),
                payload.get("summary"),
                payload.get("tokens_in"),
                payload.get("tokens_out"),
            ),
        )

        # 예: 20스텝마다 WAL 체크포인트
        if step_no % 20 == 0:
            conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")

def set_state(run_id: str, key: str, value):
    now = int(time.time())
    with connect() as conn:
        conn.execute(
            "INSERT OR REPLACE INTO kv_state(run_id, key, value_json, updated_at) VALUES(?, ?, ?, ?)",
            (run_id, key, json.dumps(value, ensure_ascii=False), now),
        )

def get_state(run_id: str, key: str, default=None):
    with connect() as conn:
        row = conn.execute(
            "SELECT value_json FROM kv_state WHERE run_id=? AND key=?",
            (run_id, key),
        ).fetchone()
    if not row:
        return default
    return json.loads(row[0])

if __name__ == "__main__":
    init_schema()
    run_id = create_run()

    for i in range(1, 101):
        # 여기서는 예시로 길이가 커질 수 있는 tool_output을 일부러 생성
        tool_output = "x" * (i * 1000)
        save_step(
            run_id,
            i,
            {
                "user_input": f"step {i}",
                "model_output": "ok",
                "tool_name": "dummy",
                "tool_input": "-",
                "tool_output": tool_output,
                "summary": "short",
                "tokens_in": 100,
                "tokens_out": 50,
            },
        )
        set_state(run_id, "last_step", i)

SELECT step_no, user_input, model_output, tool_name, tool_output, summary
FROM steps
WHERE run_id = ?
ORDER BY step_no DESC
LIMIT ?;

DELETE FROM steps
WHERE run_id = ?
  AND step_no <= (
    SELECT MAX(step_no) - 5000 FROM steps WHERE run_id = ?
  );

DELETE FROM steps
WHERE run_id = ?
  AND created_at < ?;