File size: 1,223 Bytes
ce1e080
5038bcb
 
 
 
 
87c52f8
 
 
5038bcb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87c52f8
 
 
 
 
 
5038bcb
 
cbbf969
5038bcb
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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")