Spaces:
Sleeping
Sleeping
| """Database models for Trinka Agent Backend""" | |
| from sqlalchemy import Column, String, Integer, Text, DateTime, ForeignKey, JSON, Float | |
| from sqlalchemy.dialects.postgresql import UUID | |
| from sqlalchemy.sql import func | |
| from sqlalchemy.orm import relationship | |
| import uuid | |
| from database import Base | |
| def generate_uuid(): | |
| """Generate UUID as string""" | |
| return str(uuid.uuid4()) | |
| class Conversation(Base): | |
| """Conversation table - stores chat conversations""" | |
| __tablename__ = "conversations" | |
| id = Column(UUID(as_uuid=False), primary_key=True, default=generate_uuid) | |
| user_id = Column(UUID(as_uuid=False), nullable=False, index=True) | |
| document_id = Column(UUID(as_uuid=False), nullable=True, index=True) | |
| created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False) | |
| updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False) | |
| # Relationships | |
| messages = relationship("Message", back_populates="conversation", cascade="all, delete-orphan") | |
| actions = relationship("AgentAction", back_populates="conversation", cascade="all, delete-orphan") | |
| alert_groups = relationship("AlertGroupDB", back_populates="conversation", cascade="all, delete-orphan") | |
| def __repr__(self): | |
| return f"<Conversation(id={self.id}, user_id={self.user_id})>" | |
| class Message(Base): | |
| """Message table - stores conversation messages""" | |
| __tablename__ = "messages" | |
| id = Column(UUID(as_uuid=False), primary_key=True, default=generate_uuid) | |
| conversation_id = Column(UUID(as_uuid=False), ForeignKey("conversations.id", ondelete="CASCADE"), nullable=False, index=True) | |
| type = Column(String(10), nullable=False) # 'user' or 'agent' | |
| content = Column(Text, nullable=False) | |
| timestamp = Column(DateTime(timezone=True), server_default=func.now(), nullable=False) | |
| metadata = Column(JSON, nullable=True) # Store buttons, alert groups, etc. | |
| # Relationships | |
| conversation = relationship("Conversation", back_populates="messages") | |
| def __repr__(self): | |
| return f"<Message(id={self.id}, type={self.type}, conversation_id={self.conversation_id})>" | |
| class AgentAction(Base): | |
| """Agent Actions table - stores actions executed by the agent""" | |
| __tablename__ = "agent_actions" | |
| id = Column(UUID(as_uuid=False), primary_key=True, default=generate_uuid) | |
| conversation_id = Column(UUID(as_uuid=False), ForeignKey("conversations.id", ondelete="CASCADE"), nullable=False, index=True) | |
| action_type = Column(String(50), nullable=False) # 'fix_all', 'accept_suggestion', etc. | |
| input_text = Column(Text, nullable=True) | |
| output_text = Column(Text, nullable=True) | |
| changes_applied = Column(JSON, nullable=True) # Array of changes | |
| timestamp = Column(DateTime(timezone=True), server_default=func.now(), nullable=False) | |
| # Relationships | |
| conversation = relationship("Conversation", back_populates="actions") | |
| def __repr__(self): | |
| return f"<AgentAction(id={self.id}, action_type={self.action_type})>" | |
| class AlertGroupDB(Base): | |
| """Alert Groups table - stores categorized alert groups""" | |
| __tablename__ = "alert_groups" | |
| id = Column(UUID(as_uuid=False), primary_key=True, default=generate_uuid) | |
| conversation_id = Column(UUID(as_uuid=False), ForeignKey("conversations.id", ondelete="CASCADE"), nullable=False, index=True) | |
| category = Column(String(100), nullable=False) # e.g., "capitalization_spacing" | |
| tool_category = Column(String(50), nullable=False) # 'grammar', 'paraphraser', etc. | |
| alert_count = Column(Integer, default=0) | |
| alerts = Column(JSON, nullable=True) # Array of alert objects | |
| created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False) | |
| # Relationships | |
| conversation = relationship("Conversation", back_populates="alert_groups") | |
| def __repr__(self): | |
| return f"<AlertGroupDB(id={self.id}, category={self.category}, count={self.alert_count})>" | |
| # Optional: User table (if you want to track users) | |
| class User(Base): | |
| """User table - stores user information""" | |
| __tablename__ = "users" | |
| id = Column(UUID(as_uuid=False), primary_key=True, default=generate_uuid) | |
| email = Column(String(255), unique=True, nullable=True, index=True) | |
| username = Column(String(100), nullable=True) | |
| created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False) | |
| updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False) | |
| def __repr__(self): | |
| return f"<User(id={self.id}, email={self.email})>" | |
| # Optional: Document table (if you want to track documents) | |
| class Document(Base): | |
| """Document table - stores document information""" | |
| __tablename__ = "documents" | |
| id = Column(UUID(as_uuid=False), primary_key=True, default=generate_uuid) | |
| user_id = Column(UUID(as_uuid=False), ForeignKey("users.id", ondelete="CASCADE"), nullable=True, index=True) | |
| title = Column(String(255), nullable=True) | |
| content = Column(Text, nullable=True) | |
| word_count = Column(Integer, default=0) | |
| created_at = Column(DateTime(timezone=True), server_default=func.now(), nullable=False) | |
| updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False) | |
| def __repr__(self): | |
| return f"<Document(id={self.id}, title={self.title})>" | |