40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from typing import AsyncIterator
|
|
from fastapi import Request
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
|
from sqlalchemy import String
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
# User Thread Mapping Table
|
|
class UserThreadMapping(Base):
|
|
"""Mapping of users to their chat threads.
|
|
|
|
Used to keep track of which user owns which chat thread.
|
|
Also stores meta data like thread name.
|
|
"""
|
|
|
|
__tablename__ = "userThreads"
|
|
id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
|
|
userId: Mapped[int] = mapped_column(nullable=False)
|
|
threadId: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
|
|
threadName: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
|
|
|
|
# Dependency that pulls the sessionmaker off app.state
|
|
# This is set in app.py on startup in @asynccontextmanager
|
|
async def get_session(request: Request) -> AsyncIterator[AsyncSession]:
|
|
SessionLocal: async_sessionmaker[AsyncSession] = (
|
|
request.app.state.checkpoint_sessionmaker
|
|
)
|
|
async with SessionLocal() as session:
|
|
yield session
|
|
|
|
|
|
# Optional helper to init tables at startup (demo only)
|
|
async def init_models(engine) -> None:
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|