Newer
Older
cortex-hub / ai-hub / app / db / models.py
from datetime import datetime
from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey, Boolean, JSON
from sqlalchemy.orm import relationship

# The declarative_base class is already defined in database.py.
# We will import it from there to ensure all models use the same base.
from .database import Base

# --- SQLAlchemy Models ---
# These classes define the structure of the database tables and how they relate.

class Session(Base):
    """
    SQLAlchemy model for the 'sessions' table.

    Each session represents a single conversation between a user and the AI.
    It links a user to a series of messages.
    """
    __tablename__ = 'sessions'

    # Primary key for the session.
    id = Column(Integer, primary_key=True, index=True)
    # The ID of the user who owns this session.
    user_id = Column(String, index=True, nullable=False)
    # A title for the conversation, which can be generated by the AI.
    title = Column(String, index=True, nullable=True)
    # The name of the LLM model used for this session (e.g., "Gemini", "DeepSeek").
    model_name = Column(String, nullable=True)
    # Timestamp for when the session was created.
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    # Flag to indicate if the session has been archived or soft-deleted.
    is_archived = Column(Boolean, default=False, nullable=False)

    # Defines a one-to-many relationship with the Message table.
    # 'back_populates' tells SQLAlchemy that there's a corresponding relationship
    # on the other side. 'cascade' ensures that when a session is deleted,
    # all its associated messages are also deleted.
    messages = relationship("Message", back_populates="session", cascade="all, delete-orphan")
    
    def __repr__(self):
        """
        Provides a helpful string representation of the object for debugging.
        """
        return f"<Session(id={self.id}, title='{self.title}', user_id='{self.user_id}')>"


class Message(Base):
    """
    SQLAlchemy model for the 'messages' table.

    This table stores the individual chat messages within a session,
    including who sent them (user or AI) and the content.
    """
    __tablename__ = 'messages'

    # Primary key for the message.
    id = Column(Integer, primary_key=True, index=True)
    # The foreign key that links this message to its parent session.
    # This is a critical link for reconstructing chat history.
    session_id = Column(Integer, ForeignKey('sessions.id'), nullable=False)
    # Identifies the sender of the message, e.g., 'user' or 'assistant'.
    sender = Column(String, nullable=False)
    # The actual text content of the message.
    content = Column(Text, nullable=False)
    # Timestamp for when the message was sent.
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    # The time taken for the model to generate the response, in seconds.
    model_response_time = Column(Integer, nullable=True)
    # The number of tokens in the message (both input and output).
    token_count = Column(Integer, nullable=True)
    # A JSON field to store unstructured metadata about the message, such as tool calls.
    # This column has been renamed from 'metadata' to avoid a conflict.
    message_metadata = Column(JSON, nullable=True)


    # Relationship back to the parent Session.
    # This allows us to access the parent Session object from a Message object.
    session = relationship("Session", back_populates="messages")

    def __repr__(self):
        """
        Provides a helpful string representation of the object for debugging.
        """
        return f"<Message(id={self.id}, session_id={self.session_id}, sender='{self.sender}')>"


class Document(Base):
    """
    SQLAlchemy model for the 'documents' table.

    This table stores the metadata and original text content of a document.
    The content is the data that will be chunked, embedded, and used for RAG.
    """
    __tablename__ = 'documents'

    # Primary key for the document, uniquely identifying each entry.
    id = Column(Integer, primary_key=True, index=True)
    # The title of the document for easy human-readable reference.
    title = Column(String, index=True, nullable=False)
    # The actual text content of the document. Using Text for potentially long strings.
    text = Column(Text, nullable=False)
    # The original source URL or path of the document.
    source_url = Column(String, nullable=True)
    # A string to identify the author of the document.
    author = Column(String, nullable=True)
    # The current processing status of the document (e.g., 'ready', 'processing', 'failed').
    status = Column(String, default="processing", nullable=False)
    # Timestamp for when the document was added to the database.
    created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
    # A string to identify the user who added the document, useful for multi-user apps.
    user_id = Column(String, index=True, nullable=True)
    
    # Defines a one-to-one relationship with the VectorMetadata table.
    vector_metadata = relationship(
        "VectorMetadata",
        back_populates="document",
        cascade="all, delete-orphan", # Deletes vector metadata when the document is deleted.
        uselist=False
    )

    def __repr__(self):
        """
        Provides a helpful string representation of the object for debugging.
        """
        return f"<Document(id={self.id}, title='{self.title}', user_id='{self.user_id}')>"


class VectorMetadata(Base):
    """
    SQLAlchemy model for the 'vector_metadata' table.

    This table links a document to its corresponding vector representation
    in the FAISS index. It is critical for syncing data between the
    relational database and the vector store.
    """
    __tablename__ = 'vector_metadata'

    # Primary key for the metadata entry.
    id = Column(Integer, primary_key=True, index=True)
    # Foreign key that links this metadata entry back to its Document.
    document_id = Column(Integer, ForeignKey('documents.id'), unique=True)
    # The index number in the FAISS vector store where the vector for this document is stored.
    faiss_index = Column(Integer, nullable=False, index=True)
    # Foreign key to link this vector metadata to a specific session.
    # This is crucial for retrieving relevant RAG context for a given conversation.
    session_id = Column(Integer, ForeignKey('sessions.id'), nullable=True)
    # The name of the embedding model used to create the vector.
    embedding_model = Column(String, nullable=False)

    # Defines a many-to-one relationship with the Document table.
    document = relationship("Document", back_populates="vector_metadata")
    # Defines a many-to-one relationship with the Session table.
    session = relationship("Session")

    def __repr__(self):
        """
        Provides a helpful string representation of the object for debugging.
        """
        return f"<VectorMetadata(id={self.id}, document_id={self.document_id}, faiss_index={self.faiss_index}, session_id={self.session_id})>"