92 lines
3.3 KiB
Python
92 lines
3.3 KiB
Python
"""
|
|
Notebook Mode Models for NotebookLlama Integration
|
|
Handles file uploads and isolated analysis sessions
|
|
"""
|
|
from sqlalchemy import Column, String, Boolean, DateTime, BigInteger, Integer, ForeignKey
|
|
from sqlalchemy.dialects.postgresql import UUID
|
|
from sqlalchemy.orm import relationship
|
|
from app.database import Base
|
|
import uuid
|
|
from datetime import datetime
|
|
import enum
|
|
|
|
|
|
class ProcessingStatus(enum.Enum):
|
|
"""Processing status for document uploads"""
|
|
QUEUED = "queued"
|
|
PROCESSING = "processing"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
|
|
|
|
class NotebookSession(Base):
|
|
"""
|
|
Notebook session table
|
|
Maps our internal session to external NotebookLlama notebook
|
|
"""
|
|
__tablename__ = "notebook_sessions"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
user_id = Column(UUID(as_uuid=True), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True)
|
|
conversation_id = Column(UUID(as_uuid=True), ForeignKey("conversations.id", ondelete="CASCADE"), nullable=False, index=True)
|
|
title = Column(String(255), nullable=True)
|
|
|
|
# External NotebookLlama notebook ID (INTEGER from their API)
|
|
notebookllama_notebook_id = Column(Integer, nullable=True, unique=True)
|
|
|
|
# Pin feature: pinned sessions never expire
|
|
is_pinned = Column(Boolean, default=False, nullable=False)
|
|
|
|
# Track total file size for quota management
|
|
total_file_size = Column(BigInteger, default=0, nullable=False)
|
|
|
|
# Expiration: NULL if pinned, NOW() + 24h otherwise
|
|
expires_at = Column(DateTime, nullable=True)
|
|
|
|
created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
|
|
|
|
# Relationships
|
|
user = relationship("User", back_populates="notebook_sessions")
|
|
conversation = relationship("Conversation", back_populates="notebook_session")
|
|
uploaded_files = relationship("UploadedFile", back_populates="session", cascade="all, delete-orphan")
|
|
|
|
def __repr__(self):
|
|
return f"<NotebookSession {self.id} (Notebook: {self.notebookllama_notebook_id})>"
|
|
|
|
|
|
class UploadedFile(Base):
|
|
"""
|
|
Uploaded files table
|
|
Tracks files uploaded to NotebookLlama sessions with processing status
|
|
"""
|
|
__tablename__ = "uploaded_files"
|
|
|
|
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
|
|
session_id = Column(UUID(as_uuid=True), ForeignKey("notebook_sessions.id", ondelete="CASCADE"), nullable=False, index=True)
|
|
|
|
file_name = Column(String(255), nullable=False)
|
|
file_size = Column(BigInteger, nullable=False)
|
|
file_type = Column(String(50), nullable=True) # pdf, docx, xlsx, etc.
|
|
|
|
# Local storage path
|
|
storage_path = Column(String(500), nullable=False)
|
|
|
|
# External NotebookLlama task ID (from background processing)
|
|
notebookllama_task_id = Column(Integer, nullable=True)
|
|
|
|
# Processing status tracking
|
|
processing_status = Column(
|
|
String(50),
|
|
default=ProcessingStatus.QUEUED.value,
|
|
nullable=False
|
|
)
|
|
processing_error = Column(String(500), nullable=True)
|
|
|
|
uploaded_at = Column(DateTime, default=datetime.utcnow, nullable=False)
|
|
processed_at = Column(DateTime, nullable=True)
|
|
|
|
# Relationships
|
|
session = relationship("NotebookSession", back_populates="uploaded_files")
|
|
|
|
def __repr__(self):
|
|
return f"<UploadedFile {self.file_name} (Status: {self.processing_status})>"
|