MemoryBear/api/app/models/prompt_optimizer_model.py

import datetime
import uuid
from enum import StrEnum

from sqlalchemy import Column, ForeignKey, Text, DateTime, String, Index
from sqlalchemy.dialects.postgresql import UUID

from app.db import Base


class RoleType(StrEnum):
    """
    Enumeration of message roles used in prompt optimization conversations.

    This enum standardizes the role identifiers for messages stored in the
    prompt optimization session history, ensuring consistency across
    system-generated messages, user inputs, and assistant responses.

    Attributes:
        SYSTEM (str): Represents system-level instructions or prompts that
            define the behavior or constraints of the assistant.
        USER (str): Represents messages originating from the end user.
        ASSISTANT (str): Represents messages generated by the AI assistant.
    """
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"


class PromptOptimizerModelConfig(Base):
    """
    Prompt Optimization Model Configuration.

    This table stores system-level prompt configurations for each tenant.
    The configuration defines the base system prompt used during prompt
    optimization sessions and serves as a foundational instruction set
    for the optimization process.

    Each tenant may have one or more model configurations depending on
    business requirements.

    Table Name:
        prompt_model_config

    Columns:
        id (UUID):
            Primary key. Unique identifier for the prompt model configuration.
        tenant_id (UUID):
            Foreign key referencing `tenants.id`.
            Identifies the tenant that owns this configuration.
        system_prompt (Text):
            The system-level prompt used to guide prompt optimization logic.
        created_at (DateTime):
            Timestamp indicating when the configuration was created.
        updated_at (DateTime):
            Timestamp indicating the last update time of the configuration.

    Usage:
        - Loaded when initializing a prompt optimization session
        - Acts as the root system instruction for all subsequent prompts
    """
    __tablename__ = "prompt_model_config"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True)
    tenant_id = Column(UUID(as_uuid=True), ForeignKey("tenants.id"), nullable=False, comment="Tenant ID")
    # model_id = Column(UUID(as_uuid=True), nullable=False, comment="Model ID")
    system_prompt = Column(Text, nullable=False, comment="System Prompt")

    created_at = Column(DateTime, default=datetime.datetime.now, comment="Creation Time")
    updated_at = Column(DateTime, default=datetime.datetime.now, onupdate=datetime.datetime.now, comment="Update Time")


class PromptOptimizerSession(Base):
    """
    Prompt Optimization Session Registry.

    This table records high-level metadata for prompt optimization sessions.
    Each record represents a single logical session initiated by a user
    under a specific tenant.

    The session acts as a container for multiple conversation messages
    stored in the session history table.

    Table Name:
        prompt_opt_session_list

    Columns:
        id (UUID):
            Primary key. Internal unique identifier for the session record.
        tenant_id (UUID):
            Foreign key referencing `tenants.id`.
            Identifies the tenant under which the session is created.
        session_id (UUID):
            Public-facing session identifier used to group conversation history.
        user_id (UUID):
            Foreign key referencing `users.id`.
            Identifies the user who initiated the session.
        created_at (DateTime):
            Timestamp indicating when the session was created.

    Design Notes:
        - This table intentionally does not store message content
        - Message-level data is stored in `prompt_opt_session_history`
        - Enables efficient session listing and pagination
    """
    __tablename__ = "prompt_opt_session_list"

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True)
    tenant_id = Column(UUID(as_uuid=True), ForeignKey("tenants.id"), nullable=False, comment="Tenant ID")
    # app_id = Column(UUID(as_uuid=True), ForeignKey("apps.id"), nullable=False, comment="Application ID")
    session_id = Column(UUID(as_uuid=True), nullable=False, comment="Session ID")
    user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False, comment="User ID")

    created_at = Column(DateTime, default=datetime.datetime.now, comment="Creation Time", index=True)


class PromptOptimizerSessionHistory(Base):
    """
    Prompt Optimization Session Message History.

    This table stores the complete conversational history of a prompt
    optimization session, including system prompts, user inputs, and
    assistant responses.

    Each record represents a single message within a session, preserving
    the chronological order of interactions.

    Table Name:
        prompt_opt_session_history

    Columns:
        id (UUID):
            Primary key. Unique identifier for the message record.
        tenant_id (UUID):
            Foreign key referencing `tenants.id`.
            Identifies the tenant under which the session operates.
        session_id (UUID):
            Logical session identifier linking messages to a session.
        user_id (UUID):
            Foreign key referencing `users.id`.
            Identifies the user associated with the session.
        message_role (Text):
            Role of the message sender (e.g., system, user, assistant).
        message_content (Text):
            Raw message content generated or provided during the session.
        prompt (Text):
            The prompt snapshot used at the time of message generation.
        created_at (DateTime):
            Timestamp indicating when the message was created.

    Design Notes:
        - Supports full conversation replay and audit
        - Enables prompt evolution tracking over time
        - Indexed by creation time for efficient chronological queries
    """
    __tablename__ = "prompt_opt_session_history"

    __table_args__ = (
        Index(
            "ix_prompt_opt_session_history_session_user_created",
            "session_id",
            "user_id",
            "created_at"
        ),
    )

    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, index=True)
    tenant_id = Column(UUID(as_uuid=True), ForeignKey("tenants.id"), nullable=False, comment="Tenant ID")
    # app_id = Column(UUID(as_uuid=True), ForeignKey("apps.id"), nullable=False, comment="Application ID")
    session_id = Column(UUID(as_uuid=True), nullable=False, comment="Session ID")
    user_id = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False, comment="User ID")
    role = Column(String, nullable=False, comment="Message Role")
    content = Column(Text, nullable=False, comment="Message Content")
    # prompt = Column(Text, nullable=False, comment="Prompt")

    created_at = Column(DateTime, default=datetime.datetime.now, comment="Creation Time", index=True)