152 lines
4.5 KiB
Python
Executable File
152 lines
4.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import sys
|
|
import mysql.connector
|
|
from mysql.connector import errorcode
|
|
|
|
CFG = {
|
|
"host": os.getenv("DB_HOST", "127.0.0.1"),
|
|
"user": os.getenv("DB_USER", "cs348"),
|
|
"password": os.getenv("DB_PASS", "Cs348Group!"),
|
|
"database": os.getenv("DB_NAME", "cs348db"),
|
|
"charset": "utf8mb4",
|
|
"collation": "utf8mb4_unicode_ci",
|
|
"auth_plugin": "caching_sha2_password",
|
|
}
|
|
|
|
DDL = """
|
|
CREATE TABLE IF NOT EXISTS `ChatHistory` (
|
|
`text_id` INT NOT NULL,
|
|
`type` VARCHAR(64) NOT NULL,
|
|
`text` MEDIUMTEXT NOT NULL,
|
|
PRIMARY KEY (`text_id`)
|
|
) ENGINE=InnoDB
|
|
DEFAULT CHARSET = utf8mb4
|
|
COLLATE = utf8mb4_unicode_ci;
|
|
"""
|
|
|
|
def get_connection():
|
|
try:
|
|
return mysql.connector.connect(**CFG)
|
|
except mysql.connector.Error as e:
|
|
print(f"[fatal] MySQL connect error: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
def ensure_table(conn):
|
|
with conn.cursor() as cur:
|
|
cur.execute(DDL)
|
|
conn.commit()
|
|
|
|
def get_next_id(conn):
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT COALESCE(MAX(text_id), 0) FROM `ChatHistory`;")
|
|
(max_id,) = cur.fetchone()
|
|
return int(max_id) + 1
|
|
|
|
def insert_two_rows(conn, user_text: str):
|
|
nxt = get_next_id(conn)
|
|
rows = [
|
|
(nxt, "user", user_text),
|
|
(nxt + 1, "assistant", "LLM answer here"),
|
|
]
|
|
with conn.cursor() as cur:
|
|
cur.executemany(
|
|
"INSERT INTO `ChatHistory` (`text_id`, `type`, `text`) VALUES (%s, %s, %s);",
|
|
rows,
|
|
)
|
|
conn.commit()
|
|
return nxt, nxt + 1
|
|
|
|
def list_history(conn, limit: int | None = None):
|
|
sql = "SELECT `text_id`, `type`, `text` FROM `ChatHistory` ORDER BY `text_id` ASC"
|
|
params = ()
|
|
if limit is not None and limit > 0:
|
|
sql += " LIMIT %s"
|
|
params = (limit,)
|
|
with conn.cursor() as cur:
|
|
cur.execute(sql, params)
|
|
rows = cur.fetchall()
|
|
if not rows:
|
|
print("(no rows)")
|
|
return
|
|
# Simple pretty print without extra deps
|
|
idw = max(len("text_id"), *(len(str(r[0])) for r in rows))
|
|
typew = max(len("type"), *(len(str(r[1])) for r in rows))
|
|
print(f"{'text_id'.ljust(idw)} {'type'.ljust(typew)} text")
|
|
print(f"{'-'*idw} {'-'*typew} {'-'*4}")
|
|
for tid, typ, txt in rows:
|
|
print(f"{str(tid).ljust(idw)} {str(typ).ljust(typew)} {txt}")
|
|
|
|
def print_help():
|
|
print(r"""
|
|
Commands:
|
|
\list Print all history (ordered by text_id)
|
|
\list N Print first N rows (ordered by text_id)
|
|
\q Quit
|
|
\help This help
|
|
|
|
Anything else is sent as a user prompt:
|
|
- Inserts (next_id, 'user', <your text>)
|
|
- Inserts (next_id+1, 'assistant', 'LLM answer here')
|
|
""".strip())
|
|
|
|
def main():
|
|
print("Connecting to MySQL…")
|
|
conn = get_connection()
|
|
try:
|
|
ensure_table(conn)
|
|
print("Ready. Type \\help for commands. Empty line quits.")
|
|
while True:
|
|
try:
|
|
line = input("Enter prompt> ")
|
|
except (EOFError, KeyboardInterrupt):
|
|
print("\nBye.")
|
|
break
|
|
|
|
if line is None:
|
|
continue
|
|
line = line.strip()
|
|
if not line:
|
|
print("Bye.")
|
|
break
|
|
|
|
# Commands
|
|
if line == r"\q":
|
|
print("Bye.")
|
|
break
|
|
if line == r"\help":
|
|
print_help()
|
|
continue
|
|
if line.startswith(r"\list"):
|
|
parts = line.split()
|
|
limit = None
|
|
if len(parts) == 2 and parts[1].isdigit():
|
|
limit = int(parts[1])
|
|
try:
|
|
list_history(conn, limit)
|
|
except mysql.connector.Error as e:
|
|
print(f"[error] list failed: {e}", file=sys.stderr)
|
|
continue
|
|
|
|
# Regular insert path
|
|
try:
|
|
first_id, second_id = insert_two_rows(conn, line)
|
|
print(f"Inserted prompt at text_id={first_id} and placeholder at text_id={second_id}.")
|
|
except mysql.connector.Error as e:
|
|
print(f"[error] insert failed: {e}", file=sys.stderr)
|
|
if e.errno in (errorcode.CR_SERVER_GONE_ERROR, errorcode.CR_SERVER_LOST):
|
|
try:
|
|
conn.close()
|
|
except Exception:
|
|
pass
|
|
conn = get_connection()
|
|
finally:
|
|
try:
|
|
conn.close()
|
|
except Exception:
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
main()
|