From 4b5b3aafe1bed430a4821c583fb253bc61e467d8 Mon Sep 17 00:00:00 2001 From: Daniel O'Connell Date: Fri, 1 Sep 2023 20:05:03 +0200 Subject: [PATCH] MySQL support --- readme.md => README.md | 9 ++ api/.env.example | 5 + api/Pipfile | 3 + api/Pipfile.lock | 115 +++++++++++++++++- api/migrations/README | 4 + api/migrations/alembic.ini | 110 +++++++++++++++++ api/migrations/env.py | 73 +++++++++++ api/migrations/script.py.mako | 24 ++++ .../78806d965229_interactions_table.py | 37 ++++++ api/setup.py | 6 +- api/src/stampy_chat/db/__init__.py | 0 api/src/stampy_chat/db/models.py | 92 ++++++++++++++ api/src/stampy_chat/db/session.py | 82 +++++++++++++ api/src/stampy_chat/env.py | 8 ++ local_db.sh | 26 ++++ 15 files changed, 592 insertions(+), 2 deletions(-) rename readme.md => README.md (86%) create mode 100644 api/migrations/README create mode 100644 api/migrations/alembic.ini create mode 100644 api/migrations/env.py create mode 100644 api/migrations/script.py.mako create mode 100644 api/migrations/versions/78806d965229_interactions_table.py create mode 100644 api/src/stampy_chat/db/__init__.py create mode 100644 api/src/stampy_chat/db/models.py create mode 100644 api/src/stampy_chat/db/session.py create mode 100755 local_db.sh diff --git a/readme.md b/README.md similarity index 86% rename from readme.md rename to README.md index 2d8a7e0..0b8b928 100644 --- a/readme.md +++ b/README.md @@ -22,6 +22,15 @@ Discord](https://discord.com/invite/Bt8PaRTDQC) if you need help with this. Install `npm`, `python 3.11`, and [`pipenv`](https://pipenv.pypa.io/en/latest/). +### Database setup + +Some things (e.g. logging) require a database connection to work correctly. To make this easier, there is a script to set one up locally via Docker. To get this working: + +* [Install Docker](https://docs.docker.com/get-docker/) +* Run the script: `./local_db.sh` + +This should start the database, make sure it's up to date, then as a final step display a command that will allow you to connect to it directly if you want to. + ### Running a local version Open two terminal windows. In the first, run: diff --git a/api/.env.example b/api/.env.example index 4e40cd4..e22e57b 100644 --- a/api/.env.example +++ b/api/.env.example @@ -1,3 +1,8 @@ OPENAI_API_KEY="sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX" PINECONE_API_KEY="" # leave blank to use our online API instead LOGGING_URL="" # leave blank if you're not testing logging specifically + +CHAT_DB_USER="user" +CHAT_DB_PASSWORD="we all live in a yellow submarine" +CHAT_DB_HOST="127.0.0.1" +CHAT_DB_PORT="3306" diff --git a/api/Pipfile b/api/Pipfile index 41db44b..a1c6c74 100644 --- a/api/Pipfile +++ b/api/Pipfile @@ -22,6 +22,9 @@ pinecone-client = "*" python-dotenv = "*" discord-webhook = "*" requests = "*" +alembic = "*" +sqlalchemy = "*" +mysql-connector-python = "*" [dev-packages] diff --git a/api/Pipfile.lock b/api/Pipfile.lock index 3246c19..930c2d4 100644 --- a/api/Pipfile.lock +++ b/api/Pipfile.lock @@ -1,7 +1,7 @@ { "_meta": { "hash": { - "sha256": "2908b4a829b382d34a13b475862a97fc0f0839753d23aa5becb76344d0376128" + "sha256": "2b8bbf71251550eb71f393b52f3351e7a35e9ebe9bdb5d621d7f0bb2ec9bc145" }, "pipfile-spec": 6, "requires": { @@ -117,6 +117,15 @@ "markers": "python_version >= '3.7'", "version": "==1.3.1" }, + "alembic": { + "hashes": [ + "sha256:03226222f1cf943deee6c85d9464261a6c710cd19b4fe867a3ad1f25afda610f", + "sha256:8e7645c32e4f200675e69f0745415335eb59a3663f5feb487abfa0b30c45888b" + ], + "index": "pypi", + "markers": "python_version >= '3.7'", + "version": "==1.12.0" + }, "async-timeout": { "hashes": [ "sha256:4640d96be84d82d02ed59ea2b7105a0f7b33abe8703703cd0ab0bf87c427522f", @@ -374,6 +383,14 @@ "markers": "python_version >= '3.5'", "version": "==0.7.0" }, + "mako": { + "hashes": [ + "sha256:c97c79c018b9165ac9922ae4f32da095ffd3c4e6872b45eded42926deea46818", + "sha256:d60a3903dc3bb01a18ad6a89cdbe2e4eadc69c0bc8ef1e3773ba53d44c3f7a34" + ], + "markers": "python_version >= '3.7'", + "version": "==1.2.4" + }, "markupsafe": { "hashes": [ "sha256:00bc623926325b26bb9605ae9eae8a215691f33cae5df11ca5424f06f2d1f473", @@ -513,6 +530,34 @@ "markers": "python_version >= '3.7'", "version": "==6.0.4" }, + "mysql-connector-python": { + "hashes": [ + "sha256:1d292d8ee08a3e0103f3e9aba81f292c12c244c1d5345168375bf6842b2b3e9b", + "sha256:22259479bdba53e0eb64325d5d97d7ecae35e729a8b39822c16729476ba2cd92", + "sha256:40e3b5d60dbc5446ccd1901eb3db4c9b1ad07a490f4e0fc96e7ac53a6fe9b7bc", + "sha256:4f58e1d41b3ba04281099884421c3a63589204b2cf40e473c5dbebbc26971017", + "sha256:556f83cfc3e1f3e3db0080bb19ca50650815551ef7e10518053fe2ba76d2e59e", + "sha256:5f1cdf9d3ed2f3c70a91c4b94d3ab3b2890c477a61b4085bb89d2275c50e1e3d", + "sha256:64f4d7f7c60e2fe5a01d9e7915634d05a350bd866332ea4eab281b29db190e2b", + "sha256:76432810ee58abda81e0a34666f86aba69a7ded58f164a1d19835dba70b1451b", + "sha256:78a43eeb2b845986dc769ab611e3c480991c8037c94aee67ee6842d1f7e961a9", + "sha256:7abe27b88ae1354433713fc289ed663cc04f7b75ea818a2b18067c3a736b0975", + "sha256:7f1cbc06411115c318af775291386ed1494302e05b5c45da73f01593c0827bc1", + "sha256:8062c11619bc48baf9ada3f85d85832aec5e57359eb11ed9e7d75ddcce81721a", + "sha256:816b46c9dc250a0e59d11a902e187be190dae992a2b802512b888930db380487", + "sha256:89151683654748f4de59cd4cf04102b6180976923f1db9921b87cff2c2da7b1d", + "sha256:90adf52cdb25a7adc18da0cc6c5454709598af0e87343e39ab2c11d7dc5aaf34", + "sha256:9cc256eaf55a7dd706ae514bebb863d3e426bd12686b88e284bdad6c37bf6dde", + "sha256:dd0005e9b7e54c700745416ecf26435c69687d5234eef9ff11cf85a7d76945e2", + "sha256:e3a2c2b398af1c5f0cbc1b1935f98bb5a7da4f178b2e73f0851872fbb37e66cc", + "sha256:f2ac6d43fc5a01e574fcc2a6b732cc2374e6ce576ee039cc2051f18d2daafd6d", + "sha256:f36a1a308d0d0d6202fea2e51741b5265f60206be331cfbb32f2f5bf62a20359", + "sha256:f8160777788c7e561f9d7f93eff45776031c83aa2f02a0d8c138d8f9c9bd088e" + ], + "index": "pypi", + "markers": "python_version >= '3.8'", + "version": "==8.1.0" + }, "numpy": { "hashes": [ "sha256:003a9f530e880cb2cd177cba1af7220b9aa42def9c4afc2a2fc3ee6be7eb2b22", @@ -566,6 +611,26 @@ "markers": "python_version >= '3.8'", "version": "==2.2.2" }, + "protobuf": { + "hashes": [ + "sha256:1f22ac0ca65bb70a876060d96d914dae09ac98d114294f77584b0d2644fa9c30", + "sha256:237216c3326d46808a9f7c26fd1bd4b20015fb6867dc5d263a493ef9a539293b", + "sha256:27f4d15021da6d2b706ddc3860fac0a5ddaba34ab679dc182b60a8bb4e1121cc", + "sha256:299ea899484ee6f44604deb71f424234f654606b983cb496ea2a53e3c63ab791", + "sha256:3d164928ff0727d97022957c2b849250ca0e64777ee31efd7d6de2e07c494717", + "sha256:6ab80df09e3208f742c98443b6166bcb70d65f52cfeb67357d52032ea1ae9bec", + "sha256:78a28c9fa223998472886c77042e9b9afb6fe4242bd2a2a5aced88e3f4422aa7", + "sha256:7cd532c4566d0e6feafecc1059d04c7915aec8e182d1cf7adee8b24ef1e2e6ab", + "sha256:89f9149e4a0169cddfc44c74f230d7743002e3aa0b9472d8c28f0388102fc4c2", + "sha256:a53fd3f03e578553623272dc46ac2f189de23862e68565e83dde203d41b76fc5", + "sha256:b135410244ebe777db80298297a97fbb4c862c881b4403b71bac9d4107d61fd1", + "sha256:b98d0148f84e3a3c569e19f52103ca1feacdac0d2df8d6533cf983d1fda28462", + "sha256:d1736130bce8cf131ac7957fa26880ca19227d4ad68b4888b3be0dea1f95df97", + "sha256:f45460f9ee70a0ec1b6694c6e4e348ad2019275680bd68a1d9314b8c7e01e574" + ], + "markers": "python_version >= '3.7'", + "version": "==4.21.12" + }, "python-dateutil": { "hashes": [ "sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86", @@ -758,6 +823,54 @@ "markers": "python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3'", "version": "==1.16.0" }, + "sqlalchemy": { + "hashes": [ + "sha256:1506e988ebeaaf316f183da601f24eedd7452e163010ea63dbe52dc91c7fc70e", + "sha256:1a58052b5a93425f656675673ef1f7e005a3b72e3f2c91b8acca1b27ccadf5f4", + "sha256:1b74eeafaa11372627ce94e4dc88a6751b2b4d263015b3523e2b1e57291102f0", + "sha256:1be86ccea0c965a1e8cd6ccf6884b924c319fcc85765f16c69f1ae7148eba64b", + "sha256:1d35d49a972649b5080557c603110620a86aa11db350d7a7cb0f0a3f611948a0", + "sha256:243d0fb261f80a26774829bc2cee71df3222587ac789b7eaf6555c5b15651eed", + "sha256:26a3399eaf65e9ab2690c07bd5cf898b639e76903e0abad096cd609233ce5208", + "sha256:27d554ef5d12501898d88d255c54eef8414576f34672e02fe96d75908993cf53", + "sha256:3364b7066b3c7f4437dd345d47271f1251e0cfb0aba67e785343cdbdb0fff08c", + "sha256:3423dc2a3b94125094897118b52bdf4d37daf142cbcf26d48af284b763ab90e9", + "sha256:3c6aceebbc47db04f2d779db03afeaa2c73ea3f8dcd3987eb9efdb987ffa09a3", + "sha256:3ce5e81b800a8afc870bb8e0a275d81957e16f8c4b62415a7b386f29a0cb9763", + "sha256:411e7f140200c02c4b953b3dbd08351c9f9818d2bd591b56d0fa0716bd014f1e", + "sha256:4cde2e1096cbb3e62002efdb7050113aa5f01718035ba9f29f9d89c3758e7e4e", + "sha256:5768c268df78bacbde166b48be788b83dddaa2a5974b8810af422ddfe68a9bc8", + "sha256:599ccd23a7146e126be1c7632d1d47847fa9f333104d03325c4e15440fc7d927", + "sha256:5ed61e3463021763b853628aef8bc5d469fe12d95f82c74ef605049d810f3267", + "sha256:63a368231c53c93e2b67d0c5556a9836fdcd383f7e3026a39602aad775b14acf", + "sha256:63e73da7fb030ae0a46a9ffbeef7e892f5def4baf8064786d040d45c1d6d1dc5", + "sha256:6eb6d77c31e1bf4268b4d61b549c341cbff9842f8e115ba6904249c20cb78a61", + "sha256:6f8a934f9dfdf762c844e5164046a9cea25fabbc9ec865c023fe7f300f11ca4a", + "sha256:6fe7d61dc71119e21ddb0094ee994418c12f68c61b3d263ebaae50ea8399c4d4", + "sha256:759b51346aa388c2e606ee206c0bc6f15a5299f6174d1e10cadbe4530d3c7a98", + "sha256:76fdfc0f6f5341987474ff48e7a66c3cd2b8a71ddda01fa82fedb180b961630a", + "sha256:77d37c1b4e64c926fa3de23e8244b964aab92963d0f74d98cbc0783a9e04f501", + "sha256:79543f945be7a5ada9943d555cf9b1531cfea49241809dd1183701f94a748624", + "sha256:79fde625a0a55220d3624e64101ed68a059c1c1f126c74f08a42097a72ff66a9", + "sha256:7d3f175410a6db0ad96b10bfbb0a5530ecd4fcf1e2b5d83d968dd64791f810ed", + "sha256:8dd77fd6648b677d7742d2c3cc105a66e2681cc5e5fb247b88c7a7b78351cf74", + "sha256:a3f0dd6d15b6dc8b28a838a5c48ced7455c3e1fb47b89da9c79cc2090b072a50", + "sha256:bcb04441f370cbe6e37c2b8d79e4af9e4789f626c595899d94abebe8b38f9a4d", + "sha256:c3d99ba99007dab8233f635c32b5cd24fb1df8d64e17bc7df136cedbea427897", + "sha256:ca8a5ff2aa7f3ade6c498aaafce25b1eaeabe4e42b73e25519183e4566a16fc6", + "sha256:cb0d3e94c2a84215532d9bcf10229476ffd3b08f481c53754113b794afb62d14", + "sha256:d1b09ba72e4e6d341bb5bdd3564f1cea6095d4c3632e45dc69375a1dbe4e26ec", + "sha256:d32b5ffef6c5bcb452723a496bad2d4c52b346240c59b3e6dba279f6dcc06c14", + "sha256:d3793dcf5bc4d74ae1e9db15121250c2da476e1af8e45a1d9a52b1513a393459", + "sha256:dd81466bdbc82b060c3c110b2937ab65ace41dfa7b18681fdfad2f37f27acdd7", + "sha256:e4e571af672e1bb710b3cc1a9794b55bce1eae5aed41a608c0401885e3491179", + "sha256:ea8186be85da6587456c9ddc7bf480ebad1a0e6dcbad3967c4821233a4d4df57", + "sha256:eefebcc5c555803065128401a1e224a64607259b5eb907021bf9b175f315d2a6" + ], + "index": "pypi", + "markers": "python_version >= '3.7'", + "version": "==2.0.20" + }, "stampy-chat": { "editable": true, "path": "." diff --git a/api/migrations/README b/api/migrations/README new file mode 100644 index 0000000..6f0ad62 --- /dev/null +++ b/api/migrations/README @@ -0,0 +1,4 @@ +# Stampy chat database migrations + +Run `alembic upgrade head` to update the database to the newest state. +Run `alembic downgrade -` to revert the last `` migrations. \ No newline at end of file diff --git a/api/migrations/alembic.ini b/api/migrations/alembic.ini new file mode 100644 index 0000000..a0f547c --- /dev/null +++ b/api/migrations/alembic.ini @@ -0,0 +1,110 @@ +# A generic, single database configuration. + +[alembic] +# path to migration scripts +script_location = migrations + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +# Uncomment the line below if you want the files to be prepended with date and time +# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file +# for all available tokens +# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. +prepend_sys_path = . + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the python-dateutil library that can be +# installed by adding `alembic[tz]` to the pip requirements +# string value is passed to dateutil.tz.gettz() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the +# "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to migrations/versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "version_path_separator" below. +# version_locations = %(here)s/bar:%(here)s/bat:migrations/versions + +# version path separator; As mentioned above, this is the character used to split +# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep. +# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas. +# Valid values for version_path_separator are: +# +# version_path_separator = : +# version_path_separator = ; +# version_path_separator = space +version_path_separator = os # Use os.pathsep. Default configuration used for new projects. + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# new in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +sqlalchemy.url = driver://user:pass@localhost/dbname + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# Logging configuration +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/api/migrations/env.py b/api/migrations/env.py new file mode 100644 index 0000000..6ecb946 --- /dev/null +++ b/api/migrations/env.py @@ -0,0 +1,73 @@ +from logging.config import fileConfig + +from sqlalchemy import engine_from_config +from sqlalchemy import pool + +from alembic import context + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + +# Interpret the config file for Python logging. +# This line sets up loggers basically. +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +from stampy_chat.env import DB_CONNECTION_URI +config.set_main_option("sqlalchemy.url", DB_CONNECTION_URI) + +from stampy_chat.db.models import Base +target_metadata = Base.metadata + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure( + connection=connection, target_metadata=target_metadata + ) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/api/migrations/script.py.mako b/api/migrations/script.py.mako new file mode 100644 index 0000000..55df286 --- /dev/null +++ b/api/migrations/script.py.mako @@ -0,0 +1,24 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = ${repr(branch_labels)} +depends_on = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/api/migrations/versions/78806d965229_interactions_table.py b/api/migrations/versions/78806d965229_interactions_table.py new file mode 100644 index 0000000..057a630 --- /dev/null +++ b/api/migrations/versions/78806d965229_interactions_table.py @@ -0,0 +1,37 @@ +"""Interactions table + +Revision ID: 78806d965229 +Revises: +Create Date: 2023-08-24 16:26:36.265228 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import mysql +from stampy_chat.db.models import UUID + +# revision identifiers, used by Alembic. +revision = '78806d965229' +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + 'interactions', + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('session_id', UUID(length=16), nullable=False), + sa.Column('interaction_no', sa.Integer(), nullable=False), + sa.Column('query', sa.String(length=1028), nullable=False), + sa.Column('prompt', mysql.LONGTEXT(), nullable=True), + sa.Column('response', mysql.LONGTEXT(), nullable=True), + sa.Column('chunks', sa.String(length=1028), nullable=True), + sa.Column('date_created', sa.DateTime(), nullable=False), + sa.Column('moderation', sa.JSON(), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + + +def downgrade() -> None: + op.drop_table('interactions') diff --git a/api/setup.py b/api/setup.py index 0e47c7a..b83188e 100644 --- a/api/setup.py +++ b/api/setup.py @@ -20,6 +20,10 @@ setup( 'pinecone-client', 'python-dotenv', 'discord-webhook', - 'requests' + 'requests', + + 'alembic', + 'sqlalchemy', + 'mysql-connector-python', ], ) diff --git a/api/src/stampy_chat/db/__init__.py b/api/src/stampy_chat/db/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/src/stampy_chat/db/models.py b/api/src/stampy_chat/db/models.py new file mode 100644 index 0000000..06bc435 --- /dev/null +++ b/api/src/stampy_chat/db/models.py @@ -0,0 +1,92 @@ +import logging +import uuid +from datetime import datetime +from typing import Optional + +from sqlalchemy import ( + BINARY, JSON, DateTime, Integer, String, and_, func, select +) +from sqlalchemy.dialects.mysql import LONGTEXT +from sqlalchemy.ext.hybrid import hybrid_property +from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column +from sqlalchemy.types import TypeDecorator + +logger = logging.getLogger(__name__) + + +class UUID(TypeDecorator): + + impl = BINARY(16) + + def process_bind_param(self, value, dialect): + if not value: + value = uuid.uuid4() + elif not isinstance(value, uuid.UUID): + value = uuid.UUID(value) + + if dialect.name == "mysql": + return value.bytes + return value + + def process_result_value(self, value, dialect): + if value and dialect.name == "mysql": + return uuid.UUID(value.hex()) + return value + + +class Base(DeclarativeBase): + pass + + +class Interaction(Base): + __tablename__ = "interactions" + + id: Mapped[int] = mapped_column("id", primary_key=True) + + # The session_id is set once per session, so can be easily used to extract whole histories + session_id: Mapped[str] = mapped_column(UUID(), default=uuid.uuid4) + + # The 0-indexed number of the current conversation pair, where a conversation pair is a prompt + its response. + # The system prompt isn't counted as part of the interaction list + interaction_no: Mapped[int] = mapped_column(Integer) + + # The actual query provided by the user + query: Mapped[str] = mapped_column(String(1028)) + + # The full prompt as sent to the LLM + prompt: Mapped[Optional[str]] = mapped_column(LONGTEXT) + + # Whatever the LLM returns + response: Mapped[Optional[str]] = mapped_column(LONGTEXT) + + # The ids of the chunks used for the prompt + chunks: Mapped[Optional[str]] = mapped_column(String(1028)) # TODO: Change this to a proper format + date_created: Mapped[datetime] = mapped_column(DateTime, default=func.now()) + + # Any moderation data + moderation: Mapped[Optional[JSON]] = mapped_column(JSON, default="{}") + + @hybrid_property + def history(self): + return Session.object_session(self).query(Interaction).filter( + and_( + Interaction.session_id == self.session_id, + Interaction.interaction_no < self.interaction_no + ) + ) + + @history.expression + def history(cls): + # This part is for the class level expression + return ( + select([Interaction]). + where( + and_( + Interaction.session_id == cls.session_id, + Interaction.interaction_no < cls.interaction_no + ) + ) + ) + + def __repr__(self) -> str: + return f"Interaction(session={self.session_id!r}, no={self.interaction_no!r}, query={self.query!r}, response={self.response!r})" diff --git a/api/src/stampy_chat/db/session.py b/api/src/stampy_chat/db/session.py new file mode 100644 index 0000000..57ca3cb --- /dev/null +++ b/api/src/stampy_chat/db/session.py @@ -0,0 +1,82 @@ +import time +import logging +from contextlib import contextmanager +from sqlalchemy import create_engine +from sqlalchemy.orm import Session +from sqlalchemy.exc import SQLAlchemyError +from stampy_chat.env import DB_CONNECTION_URI + + +logger = logging.getLogger(__name__) + +# We create a single engine for the entire application +engine = create_engine(DB_CONNECTION_URI, echo=False) + + +@contextmanager +def make_session(auto_commit=False): + with Session(engine, autoflush=False) as session: + yield session + if auto_commit: + session.commit() + + +class ItemAdder: + """A helper class to manage adding and flushing items to the database. + + This class exposes an `add(*items)` method which will add any provided items to the + session, and if needed commit it. It will also handle rollbacks if an error occurs + while writing. + + Commits happen whenever more than `batch_size` items have been added since the last + commit, or more than `save_every` seconds have passed - whichever is first. + """ + + def __init__(self, engine=None, batch_size=100, save_every=1): + """Initialise the adder. + + :param sqlalchemy.Engine engine: The engine to be used for connections. Will create one if not provided + :param int batch_size: will commit the session once this many items have been added + :param int save_every: will commit the session if this many seconds have passed since the last addition + """ + self.engine = engine or create_engine(DB_CONNECTION_URI, echo=False) + self.batch_size = batch_size + self.save_every = save_every + + self.session = Session(self.engine) + self._counter = 0 + self._last_save = time.time() + + def commit(self): + try: + self.session.commit() + except SQLAlchemyError as e: + logger.warn('Got error when trying to commit to database: %s', e) + self.session.rollback() + raise e + self._last_save = time.time() + + @contextmanager + def get_session(self): + if not self.session: + logger.info('No session - creating new one') + self.session = Session(self.engine) + + yield self.session + + if (self._counter % self.batch_size) == 0 or time.time() - self._last_save > self.save_every: + logger.info('Commiting batch to database') + self.commit() + + def add(self, *items): + """Add the provided items to the database, commiting them if needed.""" + with self.get_session() as session: + session.add_all(items) + self._counter += len(items) + logger.debug('added %s items', len(items)) + + def __del__(self): + logger.debug('cleaning up session') + if self.session: + self.commit() + self.session.close() diff --git a/api/src/stampy_chat/env.py b/api/src/stampy_chat/env.py index 0003f47..211e4a0 100644 --- a/api/src/stampy_chat/env.py +++ b/api/src/stampy_chat/env.py @@ -44,3 +44,11 @@ def log(*args, end="\n"): m_section = "```\n" + m_section + "\n```" DiscordWebhook(url=DISCORD_LOGGING_URL, content=m_section).execute() DiscordWebhook(url=DISCORD_LOGGING_URL, content="```\n" + message + "\n```").execute() + +### MySQL ### +user = os.environ.get("CHAT_DB_USER", "user") +password = os.environ.get("CHAT_DB_PASSWORD", "we all live in a yellow submarine") +host = os.environ.get("CHAT_DB_HOST", "127.0.0.1") +port = os.environ.get("CHAT_DB_PORT", "3306") +db_name = os.environ.get("CHAT_DB_NAME", "stampy_chat") +DB_CONNECTION_URI = f"mysql+mysqlconnector://{user}:{password}@{host}:{port}/{db_name}" diff --git a/local_db.sh b/local_db.sh new file mode 100755 index 0000000..e02db95 --- /dev/null +++ b/local_db.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +ROOT_PASSWORD=my-secret-pw + +docker start stampy-db +if [ $? -ne 0 ]; then + echo 'No docker container found - creating a new one' + docker run --name stampy-db -p 3306:3306 -e MYSQL_ROOT_PASSWORD=$ROOT_PASSWORD -d mysql:latest +fi + +echo "Waiting till mysql is available..." +while ! mysql -h 127.0.0.1 --user root --password=$ROOT_PASSWORD -e "SELECT 1" ; do + sleep 5 +done + +echo "Setting up database..." +mysql -h 127.0.0.1 -u root -p$ROOT_PASSWORD << EOF +CREATE DATABASE IF NOT EXISTS stampy_chat; +CREATE USER IF NOT EXISTS user IDENTIFIED BY 'we all live in a yellow submarine'; +GRANT ALL PRIVILEGES ON stampy_chat.* TO user; +EOF + +echo "Running migrations" + +cd "$(dirname "$0")/api" && pipenv run alembic --config migrations/alembic.ini upgrade head + +echo "The database is set up. Connect to it via 'mysql -h 127.0.0.1 -u user \"--password=we all live in a yellow submarine\" stampy_chat'"