180 lines
7.4 KiB
Python
180 lines
7.4 KiB
Python
from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, ForeignKey, Boolean, Enum as SQLEnum
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
from sqlalchemy.orm import sessionmaker, relationship, Session
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
import enum
|
|
import os
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
Base = declarative_base()
|
|
|
|
# Database connection
|
|
DATABASE_URL = f"postgresql+psycopg2://{os.getenv('pgql_user')}:{os.getenv('pgql_psw')}@localhost:5432/{os.getenv('pgql_db')}"
|
|
engine = create_engine(DATABASE_URL, pool_size=10, max_overflow=20)
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
|
|
|
|
class PermissionLevel(enum.Enum):
|
|
READ = "READ"
|
|
WRITE = "WRITE"
|
|
SHARE = "SHARE" # Write + Share permissions
|
|
ADMIN = "ADMIN"
|
|
|
|
|
|
class User(Base):
|
|
__tablename__ = "users"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
email = Column(String(255), unique=True, index=True, nullable=False)
|
|
username = Column(String(100), unique=True, index=True, nullable=False)
|
|
password_hash = Column(String(255), nullable=False)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
# Relationships
|
|
documents = relationship("Document", back_populates="owner", foreign_keys="Document.user_id")
|
|
notebooks = relationship("Notebook", back_populates="owner")
|
|
chat_sessions = relationship("ChatSession", back_populates="user")
|
|
shared_with_me = relationship("DocumentShare", foreign_keys="DocumentShare.shared_with_user_id", back_populates="shared_with_user")
|
|
shared_by_me = relationship("DocumentShare", foreign_keys="DocumentShare.owner_id", back_populates="owner")
|
|
|
|
|
|
class Document(Base):
|
|
__tablename__ = "documents"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
|
|
filename = Column(String(500), nullable=False)
|
|
original_filename = Column(String(500), nullable=False)
|
|
llamacloud_file_id = Column(String(255), nullable=True)
|
|
pipeline_id = Column(String(255), nullable=True)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
# Relationships
|
|
owner = relationship("User", back_populates="documents", foreign_keys=[user_id])
|
|
document_summaries = relationship("DocumentSummary", back_populates="document")
|
|
notebook_links = relationship("NotebookDocument", back_populates="document")
|
|
shares = relationship("DocumentShare", back_populates="document")
|
|
|
|
|
|
class Notebook(Base):
|
|
"""Collection of documents (like Google NotebookLM)"""
|
|
__tablename__ = "notebooks"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
|
|
name = Column(String(500), nullable=False)
|
|
description = Column(Text, nullable=True)
|
|
podcast_path = Column(String(500), nullable=True)
|
|
podcast_generated_at = Column(DateTime, nullable=True)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
|
|
|
# Relationships
|
|
owner = relationship("User", back_populates="notebooks")
|
|
document_links = relationship("NotebookDocument", back_populates="notebook", cascade="all, delete-orphan")
|
|
chat_sessions = relationship("ChatSession", back_populates="notebook")
|
|
|
|
|
|
class NotebookDocument(Base):
|
|
"""Junction table linking notebooks to documents"""
|
|
__tablename__ = "notebook_documents"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
notebook_id = Column(Integer, ForeignKey("notebooks.id"), nullable=False, index=True)
|
|
document_id = Column(Integer, ForeignKey("documents.id"), nullable=False, index=True)
|
|
display_order = Column(Integer, default=0)
|
|
added_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
# Relationships
|
|
notebook = relationship("Notebook", back_populates="document_links")
|
|
document = relationship("Document", back_populates="notebook_links")
|
|
|
|
|
|
class DocumentSummary(Base):
|
|
"""Summary/analysis of individual document (renamed from old Notebook)"""
|
|
__tablename__ = "document_summaries"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
|
|
document_id = Column(Integer, ForeignKey("documents.id"), nullable=False, index=True)
|
|
summary = Column(Text, nullable=True)
|
|
highlights = Column(Text, nullable=True) # JSON array stored as text
|
|
questions = Column(Text, nullable=True) # JSON array stored as text
|
|
answers = Column(Text, nullable=True) # JSON array stored as text
|
|
md_content = Column(Text, nullable=True)
|
|
mind_map_path = Column(String(500), nullable=True)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
# Relationships
|
|
document = relationship("Document", back_populates="document_summaries")
|
|
|
|
|
|
class ChatSession(Base):
|
|
__tablename__ = "chat_sessions"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
|
|
document_id = Column(Integer, ForeignKey("documents.id"), nullable=True, index=True) # Made nullable
|
|
notebook_id = Column(Integer, ForeignKey("notebooks.id"), nullable=True, index=True) # NEW: for notebook-level chats
|
|
title = Column(String(500), nullable=True)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
|
|
|
# Relationships
|
|
user = relationship("User", back_populates="chat_sessions")
|
|
notebook = relationship("Notebook", back_populates="chat_sessions")
|
|
messages = relationship("ChatMessage", back_populates="session", cascade="all, delete-orphan")
|
|
|
|
|
|
class ChatMessage(Base):
|
|
__tablename__ = "chat_messages"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
session_id = Column(Integer, ForeignKey("chat_sessions.id"), nullable=False, index=True)
|
|
role = Column(String(50), nullable=False) # 'user' or 'assistant'
|
|
content = Column(Text, nullable=False)
|
|
sources = Column(Text, nullable=True)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
# Relationships
|
|
session = relationship("ChatSession", back_populates="messages")
|
|
|
|
|
|
class DocumentShare(Base):
|
|
__tablename__ = "document_shares"
|
|
|
|
id = Column(Integer, primary_key=True, index=True)
|
|
document_id = Column(Integer, ForeignKey("documents.id"), nullable=True, index=True)
|
|
notebook_id = Column(Integer, ForeignKey("notebooks.id"), nullable=True, index=True) # NEW: share notebooks too
|
|
owner_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
|
|
shared_with_user_id = Column(Integer, ForeignKey("users.id"), nullable=False, index=True)
|
|
permission_level = Column(SQLEnum(PermissionLevel), nullable=False, default=PermissionLevel.READ)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
|
|
# Relationships
|
|
document = relationship("Document", back_populates="shares")
|
|
owner = relationship("User", foreign_keys=[owner_id], back_populates="shared_by_me")
|
|
shared_with_user = relationship("User", foreign_keys=[shared_with_user_id], back_populates="shared_with_me")
|
|
|
|
|
|
# Helper functions
|
|
def get_db() -> Session:
|
|
"""Get database session"""
|
|
db = SessionLocal()
|
|
try:
|
|
return db
|
|
finally:
|
|
pass
|
|
|
|
|
|
def init_db():
|
|
"""Initialize database tables"""
|
|
Base.metadata.create_all(bind=engine)
|
|
|
|
|
|
def drop_all_tables():
|
|
"""Drop all tables (use with caution!)"""
|
|
Base.metadata.drop_all(bind=engine)
|