Spaces:
Configuration error
Configuration error
File size: 4,387 Bytes
41943e0 88f8604 41943e0 88f8604 41943e0 88f8604 |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
"""
Antigravity Notebook - Database Models
SQLAlchemy models for PostgreSQL database.
"""
from sqlalchemy import create_engine, Column, String, Integer, Text, ForeignKey, DateTime, CheckConstraint
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from sqlalchemy.dialects.postgresql import UUID, JSONB
from datetime import datetime
import uuid
from backend.config import settings
# Database engine
engine = create_engine(settings.DATABASE_URL, echo=True)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
class Notebook(Base):
"""Notebook model - container for a project"""
__tablename__ = "notebooks"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
name = Column(String(255), nullable=False)
description = Column(Text, nullable=True)
created_at = Column(DateTime, default=datetime.utcnow)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
# Relationships
sources = relationship("Source", back_populates="notebook", cascade="all, delete-orphan")
chat_messages = relationship("ChatMessage", back_populates="notebook", cascade="all, delete-orphan")
def __repr__(self):
return f"<Notebook(id={self.id}, name={self.name})>"
class Source(Base):
"""Source model - a file (PDF), URL, or text belonging to a notebook"""
__tablename__ = "sources"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
notebook_id = Column(UUID(as_uuid=True), ForeignKey("notebooks.id", ondelete="CASCADE"), nullable=False)
source_type = Column(String(10), nullable=False)
filename = Column(String(255), nullable=True)
url = Column(Text, nullable=True)
content_hash = Column(String(64), nullable=True)
created_at = Column(DateTime, default=datetime.utcnow)
source_metadata = Column(JSONB, default={})
__table_args__ = (
CheckConstraint("source_type IN ('pdf', 'url', 'text')", name="check_source_type"),
)
# Relationships
notebook = relationship("Notebook", back_populates="sources")
latent_tensors = relationship("LatentTensor", back_populates="source", cascade="all, delete-orphan")
def __repr__(self):
return f"<Source(id={self.id}, type={self.source_type}, filename={self.filename})>"
class LatentTensor(Base):
"""LatentTensor model - CLaRa-compressed representation of a segment"""
__tablename__ = "latent_tensors"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
source_id = Column(UUID(as_uuid=True), ForeignKey("sources.id", ondelete="CASCADE"), nullable=False)
tensor_path = Column(String(512), nullable=False)
segment_index = Column(Integer, nullable=False)
token_count = Column(Integer, nullable=False)
original_text_length = Column(Integer, nullable=True)
created_at = Column(DateTime, default=datetime.utcnow)
tensor_metadata = Column(JSONB, default={})
# Relationships
source = relationship("Source", back_populates="latent_tensors")
def __repr__(self):
return f"<LatentTensor(id={self.id}, source_id={self.source_id}, tokens={self.token_count})>"
class ChatMessage(Base):
"""ChatMessage model - stores chat history"""
__tablename__ = "chat_messages"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
notebook_id = Column(UUID(as_uuid=True), ForeignKey("notebooks.id", ondelete="CASCADE"), nullable=False)
role = Column(String(20), nullable=False)
content = Column(Text, nullable=False)
sources_used = Column(JSONB, default=[])
created_at = Column(DateTime, default=datetime.utcnow)
__table_args__ = (
CheckConstraint("role IN ('user', 'assistant')", name="check_role"),
)
# Relationships
notebook = relationship("Notebook", back_populates="chat_messages")
def __repr__(self):
return f"<ChatMessage(id={self.id}, role={self.role})>"
def get_db():
"""Dependency for FastAPI to get database session"""
db = SessionLocal()
try:
yield db
finally:
db.close()
def init_db():
"""Initialize database - create all tables"""
Base.metadata.create_all(bind=engine)
print("✅ Database initialized successfully!")
if __name__ == "__main__":
init_db()
|