LawBot/shared/models.py

from __future__ import annotations

from sqlalchemy import (
    BIGINT,
    BOOLEAN,
    TIMESTAMP,
    Column,
    Date,
    DateTime,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
    VARCHAR,
)
from sqlalchemy.dialects.postgresql import JSONB, TSVECTOR
from sqlalchemy.orm import declarative_base, relationship


BaseModel = declarative_base()


class User(BaseModel):
    __tablename__ = "users"

    user_id = Column(BIGINT, primary_key=True)
    username = Column(VARCHAR(33), nullable=True)
    fullname = Column(VARCHAR(128), nullable=False)
    country = Column(Text, nullable=False, default="Россия")
    region = Column(Text, nullable=True)
    user_type = Column(Text, nullable=False, default="physical_person")
    register_date = Column(DateTime(timezone=True), nullable=False)
    updated_at = Column(DateTime(timezone=True), nullable=False)


class Admin(BaseModel):
    __tablename__ = "admins"

    user_id = Column(BIGINT, primary_key=True)
    username = Column(VARCHAR(33), nullable=True)
    fullname = Column(VARCHAR(128), nullable=False)


class Blacklist(BaseModel):
    __tablename__ = "blacklist"

    user_id = Column(BIGINT, primary_key=True)


class Setting(BaseModel):
    __tablename__ = "settings"

    name = Column(String, primary_key=True)
    value = Column(JSONB, nullable=True)


class Consultation(BaseModel):
    __tablename__ = "consultations"

    id = Column(BIGINT, primary_key=True, autoincrement=True)
    user_id = Column(BIGINT, ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False)
    category = Column(Text, nullable=False)
    title = Column(Text, nullable=True)
    region = Column(Text, nullable=True)
    status = Column(Text, nullable=False, default="active")
    created_at = Column(TIMESTAMP(timezone=True), nullable=False)
    updated_at = Column(TIMESTAMP(timezone=True), nullable=False)

    messages = relationship("Message", back_populates="consultation")
    rag_queries = relationship("RagQuery", back_populates="consultation")


class Message(BaseModel):
    __tablename__ = "messages"

    id = Column(BIGINT, primary_key=True, autoincrement=True)
    consultation_id = Column(
        BIGINT,
        ForeignKey("consultations.id", ondelete="CASCADE"),
        nullable=False,
    )
    role = Column(Text, nullable=False)
    content = Column(Text, nullable=False)
    sources_json = Column(JSONB, nullable=True)
    created_at = Column(TIMESTAMP(timezone=True), nullable=False)

    consultation = relationship("Consultation", back_populates="messages")
    rag_queries = relationship("RagQuery", back_populates="user_message")


class LawSource(BaseModel):
    __tablename__ = "law_sources"

    id = Column(BIGINT, primary_key=True, autoincrement=True)
    title = Column(Text, nullable=False)
    source_type = Column(Text, nullable=False)
    jurisdiction = Column(Text, nullable=False, default="RU")
    law_type = Column(Text, nullable=True)
    document_number = Column(Text, nullable=True)
    adoption_date = Column(Date, nullable=True)
    publication_date = Column(Date, nullable=True)
    effective_date = Column(Date, nullable=True)
    source_url = Column(Text, nullable=False)
    official_publication_number = Column(Text, nullable=True)
    version_hash = Column(Text, nullable=False)
    is_active = Column(BOOLEAN, nullable=False, default=True)
    loaded_at = Column(TIMESTAMP(timezone=True), nullable=False)

    chunks = relationship("LawChunk", back_populates="source")

    __table_args__ = (
        Index("law_sources_source_url_idx", "source_url"),
        Index("law_sources_active_url_idx", "source_url", "is_active"),
    )


class LawChunk(BaseModel):
    __tablename__ = "law_chunks"

    id = Column(BIGINT, primary_key=True, autoincrement=True)
    source_id = Column(
        BIGINT,
        ForeignKey("law_sources.id", ondelete="CASCADE"),
        nullable=False,
    )
    chunk_index = Column(Integer, nullable=False)
    article_number = Column(Text, nullable=True)
    article_title = Column(Text, nullable=True)
    chunk_text = Column(Text, nullable=False)
    chunk_metadata = Column("metadata", JSONB, nullable=False)
    tsv = Column(TSVECTOR, nullable=True)
    created_at = Column(TIMESTAMP(timezone=True), nullable=False)

    source = relationship("LawSource", back_populates="chunks")

    __table_args__ = (
        Index("law_chunks_source_id_idx", "source_id"),
        Index("law_chunks_article_number_idx", "article_number"),
        Index("law_chunks_tsv_idx", "tsv", postgresql_using="gin"),
    )


class RagQuery(BaseModel):
    __tablename__ = "rag_queries"

    id = Column(BIGINT, primary_key=True, autoincrement=True)
    consultation_id = Column(
        BIGINT,
        ForeignKey("consultations.id", ondelete="CASCADE"),
        nullable=True,
    )
    user_message_id = Column(
        BIGINT,
        ForeignKey("messages.id", ondelete="CASCADE"),
        nullable=True,
    )
    generated_queries = Column(JSONB, nullable=False)
    retrieved_chunks = Column(JSONB, nullable=False)
    created_at = Column(TIMESTAMP(timezone=True), nullable=False)

    consultation = relationship("Consultation", back_populates="rag_queries")
    user_message = relationship("Message", back_populates="rag_queries")