ForeignKeyViolationError in SQLAlchemy with PostgreSQL: Steps Table Fails to Reference Threads Table

[echai-release]

Describe the bug
I am getting the following error.
2025-01-22 18:51:23,904 WARNING chainlit (Line: 95) An error occurred: (sqlalchemy.dialects.postgresql.asyncpg.IntegrityError) <class 'asyncpg.exceptions.ForeignKeyViolationError'>: insert or update on table "steps" violates foreign key constraint "steps_threadId_fkey" DETAIL: Key (threadId)=(dc26b947-dbf0-4bbb-82d5-3103a7a18437) is not present in table "threads". [SQL: INSERT INTO steps ("name", "type", "id", "threadId", "parentId", "streaming", "input", "isError", "output", "createdAt", "start", "showInput", "metadata", "generation") VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) ON CONFLICT (id) DO UPDATE SET "name" = $1, "type" = $2, "threadId" = $4, "parentId" = $5, "streaming" = $6, "input" = $7, "isError" = $8, "output" = $9, "createdAt" = $10, "start" = $11, "showInput" = $12, "metadata" = $13, "generation" = $14; ] [parameters: ('on_message', 'run', '147ffa84-749c-4f39-b5ee-7ac8787633f5', 'dc26b947-dbf0-4bbb-82d5-3103a7a18437', '4db2b8f9-cad3-4bd3-ba6f-026c513329da', False, '', False, '', '2025-01-22T18:51:22.890421Z', '2025-01-22T18:51:22.890451Z', 'json', '{}', 'null')] (Background on this error at: https://sqlalche.me/e/20/gkpj)

I am using SQLAlchemyDataLayer with Postgres database.

To Reproduce
Steps to reproduce the behavior:

1. Go to chatbot page and enter prompt then hit enter
2. See error the following error in the logs

[AidanShipperley]

I believe this was brought up in the Discord here.

Here's the solution they mentioned:

From the stacktrace, it looks like you use the legacy SQLAlchemy data layer (https://docs.chainlit.io/data-layers/sqlalchemy).
You're spot on the problem, it's the create_step in https://github.com/Chainlit/chainlit/blob/main/backend/chainlit/data/sql_alchemy.py#L347 which doesn't wait until the thread is created.

I encountered that same issue when building the Official data layer (https://docs.chainlit.io/data-layers/official). In the equivalent create_step we wait for the thread to be created before continuing with the step creation: https://github.com/Chainlit/chainlit/blob/main/backend/chainlit/data/chainlit_data_layer.py#L293

You can either drop the FK constraint in your SQLAlchemy data layer, knowing your DB state will eventually be correct.
Or we copy that bit of code from chainlit_data_layer.py to sql_alchemy.py. I don't like maintaining multiple data layers though and I know there are a few other bugs in that legacy data layer.