src/chatbot/models.py

import uuid
from datetime import datetime
from typing import List

from pgvector.sqlalchemy import Vector
from sqlalchemy import Column
from sqlmodel import Field, Relationship, SQLModel, ForeignKey

from sqlalchemy.dialects.postgresql import UUID


class KnowledgeBase(SQLModel, table=True):
    __tablename__ = "knowledge_base"
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    name: str = Field(nullable=False)
    description: str | None = None
    created_at: datetime = Field(default_factory=datetime.now)
    knowledge_chunk: List["KnowledgeChunk"] = Relationship(
        back_populates="knowledge_base"
    )


class KnowledgeChunk(SQLModel, table=True):
    __tablename__ = "knowledge_chunk"
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    kb_id: uuid.UUID = Field(
        sa_column=Column(UUID(as_uuid=True),
            ForeignKey("knowledge_base.id", ondelete="CASCADE"),
            nullable=False
        )
    )
    chunk_index: int
    chunk_text: str
    image_url: str | None = Field(default=None)
    embedding: List[float] = Field(sa_column=Column(Vector(768)))
    knowledge_base: "KnowledgeBase" = Relationship(back_populates="knowledge_chunk")