srivatsavdamaraju's picture
Upload 82 files
e947927 verified
from datetime import datetime, timezone
import enum
import uuid
from sqlalchemy import (
Column, String, Boolean, ForeignKey, Float,
DateTime, Text, Enum, Index
)
from sqlalchemy.orm import relationship
from sqlalchemy.dialects.postgresql import JSONB, UUID
import os
import sys
from pathlib import Path
# ensure project root (agent/) is on sys.path so sibling packages like "s3" can be imported
project_root = Path(__file__).resolve().parents[2] # -> ...\openai_agents\agent
if str(project_root) not in sys.path:
sys.path.insert(0, str(project_root))
from .database import Base
# ---------- ENUMS ----------
class UserRole(str, enum.Enum):
superadmin = "superadmin"
org_admin = "org_admin"
employee = "employee"
provider = "provider"
class DatasetStatus(str, enum.Enum):
pending = "pending"
approved = "approved"
rejected = "rejected"
published = "published"
class LeadStatus(str, enum.Enum):
new = "new"
contacted = "contacted"
closed = "closed"
# ---------- MODELS ----------
class Plan(Base):
__tablename__ = "plans"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
plan_type = Column(String(100), nullable=False)
price = Column(Float, nullable=False)
features = Column(Text)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
deleted_at = Column(DateTime(timezone=True), nullable=True)
created_by = Column(
UUID(as_uuid=True),
ForeignKey("users.id", ondelete="SET NULL"),
nullable=True,
index=True
)
organizations = relationship("Organization", back_populates="plan")
def __repr__(self):
return f"<Plan(id={self.id}, plan_type={self.plan_type}, price={self.price})>"
class Organization(Base):
__tablename__ = "organizations"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
name = Column(String(255), nullable=False)
domain = Column(String(255))
plan_id = Column(
UUID(as_uuid=True),
ForeignKey("plans.id", ondelete="RESTRICT"),
nullable=False,
index=True
)
location = Column(Text)
city = Column(String(100))
state = Column(String(100))
country = Column(String(100))
is_active = Column(Boolean, default=True, nullable=False)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
updated_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
onupdate=lambda: datetime.now(timezone.utc), nullable=False)
deleted_at = Column(DateTime(timezone=True), nullable=True)
created_by = Column(
UUID(as_uuid=True),
ForeignKey("users.id", ondelete="SET NULL"),
nullable=True,
index=True
)
plan = relationship("Plan", back_populates="organizations")
users = relationship("User", back_populates="organization", foreign_keys="User.organization_id")
leads = relationship("Lead", back_populates="organization")
def __repr__(self):
return f"<Organization(id={self.id}, name={self.name}, domain={self.domain})>"
class ThirdPartyProvider(Base):
__tablename__ = "third_party_providers"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
provider_name = Column(String(255), nullable=False)
is_active = Column(Boolean, default=True, nullable=False)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
deleted_at = Column(DateTime(timezone=True), nullable=True)
created_by = Column(
UUID(as_uuid=True),
ForeignKey("users.id", ondelete="SET NULL"),
nullable=True,
index=True
)
users = relationship("User", back_populates="provider", foreign_keys="User.provider_id")
datasets = relationship("Dataset", back_populates="provider")
leads = relationship("Lead", back_populates="provider")
def __repr__(self):
return f"<ThirdPartyProvider(id={self.id}, provider_name={self.provider_name})>"
class User(Base):
__tablename__ = "users"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
fullname = Column(String(255), nullable=False)
username = Column(String(100), unique=True, nullable=False, index=True)
password = Column(String(500), nullable=False) # Increased for bcrypt compatibility
email = Column(String(255), unique=True, nullable=False, index=True)
role = Column(Enum(UserRole), nullable=False, index=True)
is_active = Column(Boolean, default=True, nullable=False, index=True)
last_login = Column(DateTime(timezone=True))
organization_id = Column(
UUID(as_uuid=True),
ForeignKey("organizations.id", ondelete="SET NULL"),
nullable=True,
index=True
)
provider_id = Column(
UUID(as_uuid=True),
ForeignKey("third_party_providers.id", ondelete="SET NULL"),
nullable=True,
index=True
)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
updated_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
onupdate=lambda: datetime.now(timezone.utc), nullable=False)
deleted_at = Column(DateTime(timezone=True), nullable=True)
organization = relationship("Organization", back_populates="users", foreign_keys=[organization_id])
provider = relationship("ThirdPartyProvider", back_populates="users", foreign_keys=[provider_id])
conversation_data = relationship("ConversationData", back_populates="user", cascade="all, delete-orphan")
def __repr__(self):
return f"<User(id={self.id}, username={self.username}, email={self.email}, role={self.role})>"
class Dataset(Base):
__tablename__ = "datasets"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
dataset_name = Column(String(255), nullable=False)
dataset_type = Column(String(100))
tags = Column(String(500))
price = Column(Float)
submitted_time = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
description = Column(Text)
file_url = Column(Text)
status = Column(Enum(DatasetStatus), default=DatasetStatus.pending, nullable=False, index=True)
created_by = Column(
UUID(as_uuid=True),
ForeignKey("users.id", ondelete="SET NULL"),
nullable=True,
index=True
)
provider_id = Column(
UUID(as_uuid=True),
ForeignKey("third_party_providers.id", ondelete="RESTRICT"),
nullable=False,
index=True
)
reviewed_by = Column(
UUID(as_uuid=True),
ForeignKey("users.id", ondelete="SET NULL"),
nullable=True,
index=True
)
approved_at = Column(DateTime(timezone=True))
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
deleted_at = Column(DateTime(timezone=True), nullable=True)
provider = relationship("ThirdPartyProvider", back_populates="datasets")
reviewer = relationship("User", foreign_keys=[reviewed_by])
creator = relationship("User", foreign_keys=[created_by])
def __repr__(self):
return f"<Dataset(id={self.id}, dataset_name={self.dataset_name}, status={self.status})>"
class Lead(Base):
__tablename__ = "leads"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
organization_id = Column(
UUID(as_uuid=True),
ForeignKey("organizations.id", ondelete="CASCADE"),
nullable=False,
index=True
)
provider_id = Column(
UUID(as_uuid=True),
ForeignKey("third_party_providers.id", ondelete="RESTRICT"),
nullable=False,
index=True
)
dataset_id = Column(
UUID(as_uuid=True),
ForeignKey("datasets.id", ondelete="RESTRICT"),
nullable=False,
index=True
)
contact = Column(String(255), nullable=False)
status = Column(Enum(LeadStatus), default=LeadStatus.new, nullable=False, index=True)
note = Column(Text)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
updated_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
onupdate=lambda: datetime.now(timezone.utc), nullable=False)
deleted_at = Column(DateTime(timezone=True), nullable=True)
organization = relationship("Organization", back_populates="leads")
provider = relationship("ThirdPartyProvider", back_populates="leads")
dataset = relationship("Dataset", foreign_keys=[dataset_id])
def __repr__(self):
return f"<Lead(id={self.id}, contact={self.contact}, status={self.status})>"
class ConversationData(Base):
__tablename__ = "conversation_data"
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
convo_id = Column(String(255), nullable=False, index=True)
user_query = Column(JSONB, nullable=False)
response = Column(JSONB, nullable=False)
file_metadata = Column(JSONB, nullable=True)
is_saved = Column(Boolean, default=False, nullable=False)
user_id = Column(
UUID(as_uuid=True),
ForeignKey("users.id", ondelete="CASCADE"),
nullable=False,
index=True
)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
updated_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
onupdate=lambda: datetime.now(timezone.utc), nullable=False)
deleted_at = Column(DateTime(timezone=True), nullable=True)
user = relationship("User",back_populates="conversation_data")
def __repr__(self):
return f"<ConversationData(id={self.id}, convo_id={self.convo_id})>"\
class UserDatasetsMetadata(Base):
__tablename__ = "user_datasets_metadata"
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
)
user_metadata = Column(JSONB, nullable=False)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc), nullable=False)
updated_at = Column(
DateTime(timezone=True),
default=lambda: datetime.now(timezone.utc),
onupdate=lambda: datetime.now(timezone.utc),
nullable=False
)
# Relationship
user = relationship("User", backref="user_datasets_metadata")
def __repr__(self):
return f"<UserDatasetsMetadata(id={self.id}, user_id={self.user_id})>"