| """ |
| JARVIS — Database Layer |
| Async SQLAlchemy with SQLite. |
| Tables: sessions, conversations, embeddings_meta |
| """ |
|
|
| from datetime import datetime |
| from typing import Optional, List |
|
|
| from sqlalchemy import ( |
| Column, String, Text, DateTime, Integer, |
| ForeignKey, Index, event |
| ) |
| from sqlalchemy.ext.asyncio import ( |
| AsyncSession, create_async_engine, async_sessionmaker |
| ) |
| from sqlalchemy.orm import DeclarativeBase, relationship |
| from sqlalchemy.future import select |
| from loguru import logger |
|
|
| from app.core.config import settings |
|
|
|
|
| |
| |
| |
| engine = create_async_engine( |
| settings.DATABASE_URL, |
| echo=settings.DEBUG, |
| connect_args={"check_same_thread": False}, |
| ) |
|
|
| AsyncSessionLocal = async_sessionmaker( |
| engine, |
| class_=AsyncSession, |
| expire_on_commit=False, |
| ) |
|
|
|
|
| |
| |
| |
| class Base(DeclarativeBase): |
| pass |
|
|
|
|
| |
| |
| |
| class Session(Base): |
| """A chat session (conversation thread).""" |
| __tablename__ = "sessions" |
|
|
| id = Column(String(36), primary_key=True) |
| title = Column(String(200), default="New Chat") |
| created_at = Column(DateTime, default=datetime.utcnow) |
| updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) |
|
|
| messages = relationship("Message", back_populates="session", cascade="all, delete-orphan") |
|
|
|
|
| class Message(Base): |
| """A single message in a session.""" |
| __tablename__ = "messages" |
|
|
| id = Column(Integer, primary_key=True, autoincrement=True) |
| session_id = Column(String(36), ForeignKey("sessions.id"), nullable=False) |
| role = Column(String(20), nullable=False) |
| content = Column(Text, nullable=False) |
| modality = Column(String(20), default="text") |
| metadata_ = Column("metadata", Text, default="{}") |
| created_at = Column(DateTime, default=datetime.utcnow) |
|
|
| session = relationship("Session", back_populates="messages") |
|
|
| __table_args__ = ( |
| Index("ix_messages_session_id", "session_id"), |
| Index("ix_messages_created_at", "created_at"), |
| ) |
|
|
|
|
| class EmbeddingMeta(Base): |
| """Metadata for stored FAISS embeddings.""" |
| __tablename__ = "embeddings_meta" |
|
|
| id = Column(Integer, primary_key=True, autoincrement=True) |
| faiss_id = Column(Integer, unique=True, nullable=False) |
| session_id = Column(String(36), nullable=True) |
| text = Column(Text, nullable=False) |
| created_at = Column(DateTime, default=datetime.utcnow) |
|
|
|
|
| |
| |
| |
| async def init_db(): |
| """Create all tables if they don't exist.""" |
| async with engine.begin() as conn: |
| await conn.run_sync(Base.metadata.create_all) |
| logger.info("Database tables created/verified.") |
|
|
|
|
| |
| |
| |
| async def get_db() -> AsyncSession: |
| """FastAPI dependency — yields an async DB session.""" |
| async with AsyncSessionLocal() as session: |
| try: |
| yield session |
| await session.commit() |
| except Exception: |
| await session.rollback() |
| raise |
|
|
|
|
| |
| |
| |
| async def create_session(db: AsyncSession, session_id: str, title: str = "New Chat") -> Session: |
| sess = Session(id=session_id, title=title) |
| db.add(sess) |
| await db.flush() |
| return sess |
|
|
|
|
| async def get_session(db: AsyncSession, session_id: str) -> Optional[Session]: |
| result = await db.execute(select(Session).where(Session.id == session_id)) |
| return result.scalar_one_or_none() |
|
|
|
|
| async def list_sessions(db: AsyncSession, limit: int = 50) -> List[Session]: |
| result = await db.execute( |
| select(Session).order_by(Session.updated_at.desc()).limit(limit) |
| ) |
| return result.scalars().all() |
|
|
|
|
| async def add_message( |
| db: AsyncSession, |
| session_id: str, |
| role: str, |
| content: str, |
| modality: str = "text", |
| metadata: str = "{}", |
| ) -> Message: |
| msg = Message( |
| session_id=session_id, |
| role=role, |
| content=content, |
| modality=modality, |
| metadata_=metadata, |
| ) |
| db.add(msg) |
| await db.flush() |
| |
| sess = await get_session(db, session_id) |
| if sess: |
| sess.updated_at = datetime.utcnow() |
| return msg |
|
|
|
|
| async def get_messages( |
| db: AsyncSession, |
| session_id: str, |
| limit: int = 50, |
| ) -> List[Message]: |
| result = await db.execute( |
| select(Message) |
| .where(Message.session_id == session_id) |
| .order_by(Message.created_at.asc()) |
| .limit(limit) |
| ) |
| return result.scalars().all() |
|
|
|
|
| async def add_embedding_meta( |
| db: AsyncSession, |
| faiss_id: int, |
| text: str, |
| session_id: Optional[str] = None, |
| ) -> EmbeddingMeta: |
| em = EmbeddingMeta(faiss_id=faiss_id, text=text, session_id=session_id) |
| db.add(em) |
| await db.flush() |
| return em |
|
|
