72 lines
2.6 KiB
Python
72 lines
2.6 KiB
Python
"""RAG models: sources and chunks with pgvector embeddings."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import Any, Optional
|
|
|
|
from sqlalchemy import (
|
|
DateTime,
|
|
ForeignKey,
|
|
Integer,
|
|
JSON,
|
|
String,
|
|
Text,
|
|
UniqueConstraint,
|
|
func,
|
|
)
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
|
|
from app.core.database import Base
|
|
|
|
try:
|
|
from pgvector.sqlalchemy import Vector
|
|
|
|
EMBEDDING_DIM = 1536
|
|
_VectorType = Vector(EMBEDDING_DIM)
|
|
except Exception: # pragma: no cover
|
|
EMBEDDING_DIM = 1536
|
|
_VectorType = JSON # type: ignore[misc, assignment]
|
|
|
|
|
|
class RagSource(Base):
|
|
__tablename__ = "rag_sources"
|
|
|
|
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
|
scope: Mapped[str] = mapped_column(String(16), nullable=False, index=True) # global | project
|
|
project_id: Mapped[Optional[int]] = mapped_column(
|
|
ForeignKey("projects.id", ondelete="CASCADE"), nullable=True, index=True
|
|
)
|
|
kind: Mapped[str] = mapped_column(String(16), nullable=False) # upload | url
|
|
title: Mapped[str] = mapped_column(String(512), default="")
|
|
uri: Mapped[Optional[str]] = mapped_column(String(2048), nullable=True)
|
|
path: Mapped[Optional[str]] = mapped_column(String(2048), nullable=True)
|
|
gost_id: Mapped[Optional[str]] = mapped_column(String(128), nullable=True, index=True)
|
|
status: Mapped[str] = mapped_column(String(32), default="pending") # pending|indexing|ready|error
|
|
error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
|
|
meta: Mapped[Optional[dict]] = mapped_column(JSON, nullable=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()
|
|
)
|
|
|
|
chunks: Mapped[list[RagChunk]] = relationship(
|
|
back_populates="source", cascade="all, delete-orphan"
|
|
)
|
|
|
|
|
|
class RagChunk(Base):
|
|
__tablename__ = "rag_chunks"
|
|
__table_args__ = (UniqueConstraint("source_id", "chunk_index", name="uq_rag_chunk_source_idx"),)
|
|
|
|
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
|
source_id: Mapped[int] = mapped_column(
|
|
ForeignKey("rag_sources.id", ondelete="CASCADE"), index=True, nullable=False
|
|
)
|
|
chunk_index: Mapped[int] = mapped_column(Integer, default=0)
|
|
text: Mapped[str] = mapped_column(Text, default="")
|
|
embedding: Mapped[Optional[Any]] = mapped_column(_VectorType, nullable=True)
|
|
meta: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
|
|
|
|
source: Mapped[RagSource] = relationship(back_populates="chunks")
|