Skip to content

postgresql

PostgreSQL database implementation for chat history management.

PostgreSQLChatDatabase

Bases: BaseChatDatabase

PostgreSQLChatDatabase is a concrete implementation of BaseChatDatabase, designed to manage chat history using a PostgreSQL database.

Attributes:

Name Type Description
db_config PostgreSQLConfig

Configuration object containing database connection details.

Methods:

Name Description
init_db

Initializes the database by creating the necessary table if it doesn't exist.
add_message

Saves a message in the database.
create_session

Create a new session entry in the database.
load_session

Loads a session's messages from the database.
list_sessions

Lists all sessions in the database.
delete_session

Deletes a session from the database.
confirm_db_flush

Confirms that the user wants to flush the entire database.
flush_database

Removes all data from the database, effectively resetting it.
Source code in src/mada/core/database/postgresql.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
class PostgreSQLChatDatabase(BaseChatDatabase):
    """
    PostgreSQLChatDatabase is a concrete implementation of BaseChatDatabase,
    designed to manage chat history using a PostgreSQL database.

    Attributes:
        db_config (core.config.PostgreSQLConfig): Configuration object containing database
            connection details.

    Methods:
        init_db: Initializes the database by creating the necessary table if it
            doesn't exist.
        add_message: Saves a message in the database.
        create_session: Create a new session entry in the database.
        load_session: Loads a session's messages from the database.
        list_sessions: Lists all sessions in the database.
        delete_session: Deletes a session from the database.
        confirm_db_flush: Confirms that the user wants to flush the entire database.
        flush_database: Removes all data from the database, effectively resetting
            it.
    """

    def _connect(self):
        """
        Establish a connection to the PostgreSQL database.

        Returns:
            psycopg2.connection: Database connection object.
        """
        return psycopg2.connect(
            self.db_config.get_connection_string(), sslmode=self.db_config.sslmode
        )

    def init_db(self):
        """
        Initialize the database by creating the necessary table if it doesn't exist.
        """
        with self._connect() as conn:
            with conn.cursor() as cursor:
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS sessions (
                        session_id TEXT PRIMARY KEY,
                        last_updated TIMESTAMP
                    )
                """)
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS messages (
                        message_id SERIAL PRIMARY KEY,
                        session_id TEXT REFERENCES sessions(session_id),
                        role TEXT,
                        content TEXT,
                        timestamp TIMESTAMP
                    )
                """)
            conn.commit()

    def add_message(
        self, session_id: str, role: str, content: str, timestamp: datetime = None
    ):
        """
        Add a single message to the messages table in the database.

        Args:
            session_id (str): The ID of the session that this message is associated with
            role (str): The role (user or assistant) to designate who wrote the message
            content (str): The message contents
            timestamp (datetime): The time that the message was created
        """
        if timestamp is None:
            timestamp = datetime.now()
        with self._connect() as conn:
            with conn.cursor() as cursor:
                # Ensure the session exists
                cursor.execute(
                    """
                    INSERT INTO sessions (session_id, last_updated)
                    VALUES (%s, %s)
                    ON CONFLICT (session_id) DO NOTHING
                """,
                    (session_id, timestamp),
                )
                # Insert the message
                cursor.execute(
                    """
                    INSERT INTO messages (session_id, role, content, timestamp)
                    VALUES (%s, %s, %s, %s)
                """,
                    (session_id, role, content, timestamp),
                )
                # Update session's last_updated
                cursor.execute(
                    """
                    UPDATE sessions SET last_updated = %s WHERE session_id = %s
                """,
                    (timestamp, session_id),
                )
            conn.commit()

    def create_session(self, session_id: str):
        """
        Create a new session entry.

        Args:
            session_id (str): The ID of the session to create.
        """
        with self._connect() as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    """
                    INSERT INTO sessions (session_id, last_updated)
                    VALUES (%s, %s)
                    ON CONFLICT (session_id) DO NOTHING
                """,
                    (session_id, datetime.now()),
                )

    def load_session(self, session_id: str) -> List[dict]:
        """
        Load a session's messages from the database.

        Args:
            session_id (str): The ID of the session to load.

        Returns:
            List of messages from the session.
        """
        with self._connect() as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    """
                    SELECT role, content, timestamp FROM messages
                    WHERE session_id = %s
                    ORDER BY message_id ASC
                """,
                    (session_id,),
                )
                return [
                    {"role": row[0], "content": row[1], "timestamp": row[2]}
                    for row in cursor.fetchall()
                ]

    def list_sessions(self) -> List[Tuple[str, datetime]]:
        """
        List all sessions in the database.

        Returns:
            List of tuples containing session_id and last_updated.
        """
        with self._connect() as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    "SELECT session_id, last_updated FROM sessions ORDER BY last_updated DESC"
                )
                return [
                    (session_id, last_updated)
                    for session_id, last_updated in cursor.fetchall()
                ]

    def delete_session(self, session_id: str):
        """
        Delete a session from the database.

        Args:
            session_id (str): The ID of the session to delete.
        """
        with self._connect() as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    "DELETE FROM messages WHERE session_id = %s", (session_id,)
                )
                cursor.execute(
                    "DELETE FROM sessions WHERE session_id = %s", (session_id,)
                )
            conn.commit()

    def flush_database(self, confirm: bool = True):
        """
        Remove all data from the database, effectively resetting it.

        Warning:
            This operation is irreversible and will delete all sessions.

        Args:
            confirm (bool): If True, confirm that the user wants to flush the
                database prior to deleting anything. Otherwise, don't ask.
        """

        def _flush_db():
            LOG.info("Flushing the database...")
            with self._connect() as conn:
                with conn.cursor() as cursor:
                    cursor.execute("DELETE FROM messages")
                    cursor.execute("DELETE FROM sessions")
                conn.commit()
            LOG.info("Database successfully flushed.")

        if confirm:
            if self.confirm_db_flush():
                _flush_db()
            else:
                LOG.info("Database flush cancelled.")
        else:
            _flush_db()

add_message(session_id, role, content, timestamp=None)

Add a single message to the messages table in the database.

Parameters:

Name Type Description Default
session_id str

The ID of the session that this message is associated with

 required
role str

The role (user or assistant) to designate who wrote the message

 required
content str

The message contents

 required
timestamp datetime

The time that the message was created

 None
Source code in src/mada/core/database/postgresql.py 
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def add_message(
    self, session_id: str, role: str, content: str, timestamp: datetime = None
):
    """
    Add a single message to the messages table in the database.

    Args:
        session_id (str): The ID of the session that this message is associated with
        role (str): The role (user or assistant) to designate who wrote the message
        content (str): The message contents
        timestamp (datetime): The time that the message was created
    """
    if timestamp is None:
        timestamp = datetime.now()
    with self._connect() as conn:
        with conn.cursor() as cursor:
            # Ensure the session exists
            cursor.execute(
                """
                INSERT INTO sessions (session_id, last_updated)
                VALUES (%s, %s)
                ON CONFLICT (session_id) DO NOTHING
            """,
                (session_id, timestamp),
            )
            # Insert the message
            cursor.execute(
                """
                INSERT INTO messages (session_id, role, content, timestamp)
                VALUES (%s, %s, %s, %s)
            """,
                (session_id, role, content, timestamp),
            )
            # Update session's last_updated
            cursor.execute(
                """
                UPDATE sessions SET last_updated = %s WHERE session_id = %s
            """,
                (timestamp, session_id),
            )
        conn.commit()

create_session(session_id)

Create a new session entry.

Parameters:

Name Type Description Default
session_id str

The ID of the session to create.

 required
Source code in src/mada/core/database/postgresql.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def create_session(self, session_id: str):
    """
    Create a new session entry.

    Args:
        session_id (str): The ID of the session to create.
    """
    with self._connect() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                """
                INSERT INTO sessions (session_id, last_updated)
                VALUES (%s, %s)
                ON CONFLICT (session_id) DO NOTHING
            """,
                (session_id, datetime.now()),
            )

delete_session(session_id)

Delete a session from the database.

Parameters:

Name Type Description Default
session_id str

The ID of the session to delete.

 required
Source code in src/mada/core/database/postgresql.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def delete_session(self, session_id: str):
    """
    Delete a session from the database.

    Args:
        session_id (str): The ID of the session to delete.
    """
    with self._connect() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                "DELETE FROM messages WHERE session_id = %s", (session_id,)
            )
            cursor.execute(
                "DELETE FROM sessions WHERE session_id = %s", (session_id,)
            )
        conn.commit()

flush_database(confirm=True)

Remove all data from the database, effectively resetting it.

Warning

This operation is irreversible and will delete all sessions.

Parameters:

Name Type Description Default
confirm bool

If True, confirm that the user wants to flush the database prior to deleting anything. Otherwise, don't ask.

 True
Source code in src/mada/core/database/postgresql.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
def flush_database(self, confirm: bool = True):
    """
    Remove all data from the database, effectively resetting it.

    Warning:
        This operation is irreversible and will delete all sessions.

    Args:
        confirm (bool): If True, confirm that the user wants to flush the
            database prior to deleting anything. Otherwise, don't ask.
    """

    def _flush_db():
        LOG.info("Flushing the database...")
        with self._connect() as conn:
            with conn.cursor() as cursor:
                cursor.execute("DELETE FROM messages")
                cursor.execute("DELETE FROM sessions")
            conn.commit()
        LOG.info("Database successfully flushed.")

    if confirm:
        if self.confirm_db_flush():
            _flush_db()
        else:
            LOG.info("Database flush cancelled.")
    else:
        _flush_db()

init_db()

Initialize the database by creating the necessary table if it doesn't exist.

Source code in src/mada/core/database/postgresql.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def init_db(self):
    """
    Initialize the database by creating the necessary table if it doesn't exist.
    """
    with self._connect() as conn:
        with conn.cursor() as cursor:
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    last_updated TIMESTAMP
                )
            """)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    message_id SERIAL PRIMARY KEY,
                    session_id TEXT REFERENCES sessions(session_id),
                    role TEXT,
                    content TEXT,
                    timestamp TIMESTAMP
                )
            """)
        conn.commit()

list_sessions()

List all sessions in the database.

Returns:

Type Description
List[Tuple[str, datetime]]

List of tuples containing session_id and last_updated.
Source code in src/mada/core/database/postgresql.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def list_sessions(self) -> List[Tuple[str, datetime]]:
    """
    List all sessions in the database.

    Returns:
        List of tuples containing session_id and last_updated.
    """
    with self._connect() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                "SELECT session_id, last_updated FROM sessions ORDER BY last_updated DESC"
            )
            return [
                (session_id, last_updated)
                for session_id, last_updated in cursor.fetchall()
            ]

load_session(session_id)

Load a session's messages from the database.

Parameters:

Name Type Description Default
session_id str

The ID of the session to load.

 required

Returns:

Type Description
List[dict]

List of messages from the session.
Source code in src/mada/core/database/postgresql.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def load_session(self, session_id: str) -> List[dict]:
    """
    Load a session's messages from the database.

    Args:
        session_id (str): The ID of the session to load.

    Returns:
        List of messages from the session.
    """
    with self._connect() as conn:
        with conn.cursor() as cursor:
            cursor.execute(
                """
                SELECT role, content, timestamp FROM messages
                WHERE session_id = %s
                ORDER BY message_id ASC
            """,
                (session_id,),
            )
            return [
                {"role": row[0], "content": row[1], "timestamp": row[2]}
                for row in cursor.fetchall()
            ]