Coverage for src / quber / db / models.py: 100%

1"""SQLAlchemy models for Quber RAG database.""" 

2 

3from datetime import datetime 

4from typing import Any, Optional 

5 

6from pgvector.sqlalchemy import Vector 

7from sqlalchemy import JSON, ForeignKey, Integer, String, Text, UniqueConstraint 

8from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship 

9 

10 

11class Base(DeclarativeBase): 

12 """Base class for all database models.""" 

13 

14 pass 

15 

16 

17class Document(Base): 

18 """Represents a processed document with extracted tables.""" 

19 

20 __tablename__ = "documents" 

21 

22 id: Mapped[int] = mapped_column(Integer, primary_key=True) 

23 filename: Mapped[str] = mapped_column(String(255), nullable=False, unique=True, index=True) 

24 extraction_date: Mapped[datetime] = mapped_column(nullable=False) 

25 total_pages: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) 

26 total_tables: Mapped[Optional[int]] = mapped_column(Integer, nullable=True) 

27 executive_summary: Mapped[Optional[str]] = mapped_column(Text, nullable=True) 

28 model_provider: Mapped[Optional[str]] = mapped_column(String(50), nullable=True) 

29 model: Mapped[Optional[str]] = mapped_column(String(100), nullable=True) 

30 created_at: Mapped[datetime] = mapped_column(default=datetime.now) 

31 

32 # Relationship to extracted tables 

33 tables: Mapped[list["ExtractedTable"]] = relationship( 

34 "ExtractedTable", back_populates="document", cascade="all, delete-orphan" 

35 ) 

36 

37 def __repr__(self) -> str: 

38 return f"<Document(id={self.id}, filename='{self.filename}', tables={self.total_tables})>" 

39 

40 

41class ExtractedTable(Base): 

42 """Represents an extracted table with embeddings for semantic search.""" 

43 

44 __tablename__ = "extracted_tables" 

45 __table_args__ = (UniqueConstraint("document_id", "table_id", name="uq_document_table"),) 

46 

47 id: Mapped[int] = mapped_column(Integer, primary_key=True) 

48 document_id: Mapped[int] = mapped_column(ForeignKey("documents.id", ondelete="CASCADE"), index=True) 

49 table_id: Mapped[int] = mapped_column(Integer, nullable=False) 

50 page_number: Mapped[int] = mapped_column(Integer, nullable=False) 

51 # `extractor` records which engine produced the row — e.g. 

52 # "docling-consolidator" for the Parser+HeaderConsolidator path, or 

53 # "camelot-llm" for the CamelotLLMTableExtractor path. 

54 extractor: Mapped[Optional[str]] = mapped_column(String(50), nullable=True) 

55 procedural_title: Mapped[Optional[str]] = mapped_column(Text, nullable=True) 

56 llm_title: Mapped[Optional[str]] = mapped_column(Text, nullable=True) 

57 llm_subtitle: Mapped[Optional[str]] = mapped_column(Text, nullable=True) 

58 llm_description: Mapped[Optional[str]] = mapped_column(Text, nullable=True) 

59 table_markdown: Mapped[Optional[str]] = mapped_column(Text, nullable=True) 

60 headers: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON, nullable=True) 

61 footnotes: Mapped[Optional[list[str]]] = mapped_column(JSON, nullable=True) 

62 bbox: Mapped[Optional[list[float]]] = mapped_column(JSON, nullable=True) 

63 flavor: Mapped[Optional[str]] = mapped_column(String(20), nullable=True) 

64 table_metadata: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON, nullable=True) 

65 

66 # Vector embeddings (1024 dims for bge-large-en-v1.5) 

67 title_embedding: Mapped[Optional[Any]] = mapped_column(Vector(1024), nullable=True) 

68 description_embedding: Mapped[Optional[Any]] = mapped_column(Vector(1024), nullable=True) 

69 

70 created_at: Mapped[datetime] = mapped_column(default=datetime.now) 

71 

72 # Relationship to document 

73 document: Mapped["Document"] = relationship("Document", back_populates="tables") 

74 

75 def __repr__(self) -> str: 

76 return ( 

77 f"<ExtractedTable(id={self.id}, document_id={self.document_id}, " 

78 f"table_id={self.table_id}, page={self.page_number})>" 

79 )