memory/db_setup.py

from __future__ import annotations
import os
import psycopg2
from dotenv import load_dotenv

load_dotenv()


def get_connection():
    return psycopg2.connect(
        host     = os.getenv("DB_HOST", "127.0.0.1"),
        port     = int(os.getenv("DB_PORT", 5433)),
        dbname   = os.getenv("DB_NAME", "sre_env"),
        user     = os.getenv("DB_USER", "postgres"),
        password = os.getenv("DB_PASSWORD", ""),
    )


def setup():
    conn = get_connection()
    cur  = conn.cursor()

    # Enable pgvector extension
    cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")

    # Episodes table — stores full episode memory
    cur.execute("""
        CREATE TABLE IF NOT EXISTS episodes (
            id            SERIAL PRIMARY KEY,
            episode_id    TEXT NOT NULL,
            task_id       TEXT NOT NULL,
            task_success  BOOLEAN DEFAULT FALSE,
            total_reward  FLOAT DEFAULT 0.0,
            steps_taken   INTEGER DEFAULT 0,
            actions       JSONB DEFAULT '[]',
            summary       TEXT DEFAULT '',
            state_emb     vector(384),
            created_at    TIMESTAMP DEFAULT NOW()
        );
    """)

    # Index for fast vector similarity search
    cur.execute("""
        CREATE INDEX IF NOT EXISTS episodes_emb_idx
        ON episodes
        USING ivfflat (state_emb vector_cosine_ops)
        WITH (lists = 10);
    """)

    # Ablation records table — tracks with/without memory runs
    cur.execute("""
        CREATE TABLE IF NOT EXISTS ablation_records (
            id                       SERIAL PRIMARY KEY,
            epoch                    INTEGER,
            task_id                  TEXT,
            memory_backend           TEXT,
            mean_reward              FLOAT,
            mean_steps_to_resolution FLOAT,
            task_success_rate        FLOAT,
            created_at               TIMESTAMP DEFAULT NOW()
        );
    """)

    conn.commit()
    cur.close()
    conn.close()
    print("DB setup complete — tables and indexes created.")


if __name__ == "__main__":
    setup()