Spaces:
Sleeping
Sleeping
File size: 3,690 Bytes
58de15f |
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 |
# app/db/models.py
import uuid
from datetime import datetime
import enum as py_enum # Use alias
from sqlalchemy import (
Column, String, DateTime, ForeignKey, Text, Enum as SQLAlchemyEnum,
Integer # Import Integer type
)
# Import PostgreSQL specific types
from sqlalchemy.dialects.postgresql import UUID as PG_UUID, JSONB
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
# Import the Base class from database.py
from .database import Base
# Import Enums from schemas.py
from .schemas import MessageRole, FeedbackTypeEnum # Import both enums now
class ChatSession(Base):
"""SQLAlchemy model for chat sessions (using _v2 table)."""
__tablename__ = "chat_sessions_v2"
id = Column(PG_UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
name = Column(String, index=True, nullable=True)
created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
messages = relationship("ChatMessage", back_populates="session", cascade="all, delete-orphan", lazy="select")
def __repr__(self):
return f"<ChatSessionV2(id={self.id}, name='{self.name}')>"
class ChatMessage(Base):
"""SQLAlchemy model for chat messages (using _v2 table)."""
__tablename__ = "chat_messages_v2"
id = Column(PG_UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
session_id = Column(PG_UUID(as_uuid=True), ForeignKey("chat_sessions_v2.id"), nullable=False, index=True)
role = Column(SQLAlchemyEnum(MessageRole, name="message_role_enum_v2", create_type=False), nullable=False)
content = Column(Text, nullable=False)
created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
# --- Columns for Bot messages (Context Tracking) ---
original_query = Column(Text, nullable=True, comment="Original user query that led to this bot response")
retrieved_context_ids = Column(JSONB, nullable=True, comment="List of all chunk IDs retrieved from search for the original query")
used_context_ids = Column(JSONB, nullable=True, comment="List of chunk IDs whose content was actually used to generate this response")
attempt_number = Column(Integer, default=1, nullable=False, comment="Tracks which regeneration attempt this message represents (1 = initial, 2 = first regen, etc.)")
cumulative_used_context_ids = Column(JSONB, nullable=True, comment="Set of all context IDs used in this and previous attempts for the same original_query")
# --- END Context Tracking Columns ---
session = relationship("ChatSession", back_populates="messages")
feedback = relationship("FeedbackLog", back_populates="message", uselist=False, cascade="all, delete-orphan", lazy="select")
def __repr__(self):
attempt_str = f", attempt={self.attempt_number}" if self.role == MessageRole.BOT else ""
return f"<ChatMessageV2(id={self.id}, role='{self.role}'{attempt_str}, session_id={self.session_id})>"
class FeedbackLog(Base):
"""SQLAlchemy model for feedback logs (using _v2 table)."""
__tablename__ = "feedback_logs_v2"
id = Column(PG_UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
message_id = Column(PG_UUID(as_uuid=True), ForeignKey("chat_messages_v2.id"), nullable=False, unique=True)
feedback_type = Column(SQLAlchemyEnum(FeedbackTypeEnum, name="feedback_type_enum_v2", create_type=False), nullable=False)
feedback_comment = Column(Text, nullable=True)
created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
message = relationship("ChatMessage", back_populates="feedback")
def __repr__(self):
return f"<FeedbackLogV2(id={self.id}, message_id='{self.message_id}', type='{self.feedback_type}')>"
|