AIyaodianzhushou/backend-python/app/models/drug.py

from typing import Optional
from datetime import datetime, timezone

from sqlalchemy import Column, String, Text, Integer, Float, DateTime, JSON, func
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import Mapped, mapped_column

from app.models.user import Base


engine = create_async_engine(
    "postgresql+asyncpg://postgres:pharma2025@localhost:5432/pharmacopoeia",
    echo=False,
    pool_size=5,
    max_overflow=10,
)

async_session_factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


def get_db():
    return async_session_factory


class Drug(Base):
    __tablename__ = "drugs"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    drug_id: Mapped[str] = mapped_column(String(64), unique=True, nullable=False, index=True)
    name: Mapped[str] = mapped_column(String(256), nullable=False, index=True)
    name_en: Mapped[Optional[str]] = mapped_column(String(256))
    pinyin: Mapped[Optional[str]] = mapped_column(String(256))
    category: Mapped[Optional[str]] = mapped_column(String(64), index=True)
    subcategory: Mapped[Optional[str]] = mapped_column(String(128))
    approval_number: Mapped[Optional[str]] = mapped_column(String(64))

    sections: Mapped[Optional[dict]] = mapped_column(JSON)
    source_version: Mapped[Optional[str]] = mapped_column(String(32))
    source_volume: Mapped[Optional[str]] = mapped_column(String(32))
    source_page: Mapped[Optional[str]] = mapped_column(String(128))

    is_active: Mapped[bool] = mapped_column(default=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now()
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
    )


class DrugChunk(Base):
    __tablename__ = "drug_chunks"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    drug_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True)
    section: Mapped[Optional[str]] = mapped_column(String(64))
    content: Mapped[str] = mapped_column(Text, nullable=False)
    source: Mapped[Optional[str]] = mapped_column(Text)
    chunk_index: Mapped[int] = mapped_column(default=0)
    embedding: Mapped[Optional[list]] = mapped_column(JSON, nullable=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now()
    )