docs: add a discord rag chatbot example project

examples/discord_bot/README.md

@@ -0,0 +1,40 @@
+# Discord Bot
+
+A minimal discord bot that answers questions based on pgais documentation via RAG.
+Built with python, py-cord, sqlalchemy and pgai.
+
+## Running the Example
+
+You will need a .env file with the following variables:
+
+```shell
+OPENAI_API_KEY=xxx
+DISCORD_BOT_TOKEN=xxx
+DISCORD_CHANNEL_ID=123
+```
+
+- To get the openai api key you can create one [here](https://platform.openai.com/api-keys).
+- To create the discord bot token follow the guide by py-cord [here](https://docs.pycord.dev/en/stable/discord.html#discord-intro).
+- The channel id can be found by right clicking on the channel in discord and selecting "Copy ID". You will need to enable developer mode in discord settings to see this option.
+
+You will need to then run the vectorizer and database with:
+```
+docker compose up -d
+```
+You'll also need to install the python requirements with:
+```shell
+uv sync
+```
+
+To populate the database with our docs run the insert_docs.py script with:
+```shell
+uv run python insert_docs.py
+```
+
+Afterwards you can run the bot with:
+
+```shell
+uv run python main.py
+```
+
+The bot will answer questions by responding threads in the specified channel.

examples/discord_bot/alembic.ini

@@ -0,0 +1,39 @@
+[alembic]
+script_location = migrations
+sqlalchemy.url = postgresql+asyncpg://postgres:postgres@localhost/postgres
+revision_environment = true
+
+[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

examples/discord_bot/docker-compose.yaml

@@ -0,0 +1,19 @@
+name: pgai
+services:
+  db:
+    image: timescale/timescaledb-ha:pg17
+    environment:
+      POSTGRES_PASSWORD: postgres
+      OPENAI_API_KEY: ${OPENAI_API_KEY}
+    ports:
+      - "5432:5432"
+    volumes:
+      - data:/home/postgres/pgdata/data
+  vectorizer-worker:
+    image: timescale/pgai-vectorizer-worker:latest
+    environment:
+      PGAI_VECTORIZER_WORKER_DB_URL: postgres://postgres:postgres@db:5432/postgres
+      OPENAI_API_KEY: ${OPENAI_API_KEY}
+    command: [ "--poll-interval", "5s", "--log-level", "DEBUG" ]
+volumes:
+  data:

examples/discord_bot/migrations/README

@@ -0,0 +1 @@
+Generic single-database configuration.

examples/discord_bot/migrations/env.py

@@ -0,0 +1,92 @@
+import asyncio
+from logging.config import fileConfig
+
+from alembic.script import ScriptDirectory
+from sqlalchemy.engine import Connection
+from alembic import context
+
+# Import your models and config
+from pgai_discord_bot.main import Base
+from pgai_discord_bot.main import engine as async_engine
+
+from pgai.alembic import register_operations
+
+register_operations()
+
+# this is the Alembic Config object
+config = context.config
+
+# Interpret the config file for Python logging.
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+# Set the target metadata
+target_metadata = Base.metadata
+
+
+def process_revision_directives(context, revision, directives):
+    # extract Migration
+    migration_script = directives[0]
+    # extract current head revision
+    head_revision = ScriptDirectory.from_config(context.config).get_current_head()
+
+    if head_revision is None:
+        # edge case with first migration
+        new_rev_id = 1
+    else:
+        # default branch with incrementation
+        last_rev_id = int(head_revision.lstrip('0'))
+        new_rev_id = last_rev_id + 1
+    # fill zeros up to 4 digits: 1 -> 0001
+    migration_script.rev_id = '{0:04}'.format(new_rev_id)
+
+def include_object(object, name, type_, reflected, compare_to):
+    if type_ == "table" and name in target_metadata.info.get("pgai_managed_tables", set()):
+        return False
+    return True
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode."""
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+        process_revision_directives=process_revision_directives,
+        include_object=include_object
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+def do_run_migrations(connection: Connection) -> None:
+    context.configure(connection=connection,
+                      target_metadata=target_metadata,
+                      process_revision_directives=process_revision_directives,
+                      include_object=include_object
+                      )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+async def run_async_migrations() -> None:
+    """In this scenario we need to create an Engine
+    and associate a connection with the context."""
+
+    connectable = async_engine
+
+    async with connectable.connect() as connection:
+        await connection.run_sync(do_run_migrations)
+
+    await connectable.dispose()
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode."""
+
+    asyncio.run(run_async_migrations())
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

examples/discord_bot/migrations/script.py.mako

@@ -0,0 +1,26 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}

examples/discord_bot/migrations/versions/0001_create_documents_table.py

@@ -0,0 +1,35 @@
+"""create documents table
+
+Revision ID: 0001
+Revises: 
+Create Date: 2025-01-07 17:17:30.847468
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '0001'
+down_revision: Union[str, None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('documents',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('file_name', sa.String(), nullable=True),
+    sa.Column('content', sa.Text(), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('documents')
+    # ### end Alembic commands ###

examples/discord_bot/migrations/versions/0002_create_documents_vectorizer.py

@@ -0,0 +1,41 @@
+"""Create documents vectorizer
+
+Revision ID: 0002
+Revises: 0001
+Create Date: 2025-01-07 17:18:58.091881
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+from pgai.alembic.configuration import OpenAIConfig, RecursiveCharacterTextSplitterConfig, \
+    PythonTemplateConfig
+
+# revision identifiers, used by Alembic.
+revision: str = '0002'
+down_revision: Union[str, None] = '0001'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.execute("CREATE EXTENSION IF NOT EXISTS ai CASCADE;")
+    op.create_vectorizer(
+        source_table="documents",
+        embedding=OpenAIConfig(
+            model='text-embedding-3-small',
+            dimensions=768
+        ),
+        chunking=RecursiveCharacterTextSplitterConfig(
+            chunk_column='content',
+            chunk_size=800,
+            chunk_overlap=200,
+        ),
+        formatting=PythonTemplateConfig(template='$file_name \n $chunk')
+    )
+
+
+def downgrade() -> None:
+    op.drop_vectorizer(vectorizer_id=1, drop_all=True)