~cytrogen/evi-run

ref: 04a436312bcc8b4c23186a69f41e1f856b3d94cd evi-run/database/models.py -rw-r--r-- 4.1 KiB
04a43631 — Bendy Update README.md 6 months ago 
                                                                                
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import os

from dotenv import load_dotenv

from sqlalchemy import (
    Column, Integer, String, Boolean, DateTime, ForeignKey, Float, Text, Table, BigInteger,
    TIMESTAMP, func
)
from sqlalchemy.orm import relationship, declarative_base
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker

from config import CREDITS_USER_DAILY

load_dotenv()

# Create SQLAlchemy base
Base = declarative_base()

# Create async engine
engine = create_async_engine(
    os.getenv('DATABASE_URL'),
    echo=False,
)

# Create async session factory
async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


class User(Base):
    __tablename__ = 'users'

    telegram_id = Column(BigInteger, primary_key=True)
    language = Column(String(10), nullable=True)
    balance_credits = Column(Float, default=CREDITS_USER_DAILY)
    created_at = Column(TIMESTAMP(timezone=True), server_default=func.now())

    wallets = relationship('Wallet', back_populates='user')
    messages = relationship('ChatMessage', back_populates='user')
    payments = relationship('Payment', back_populates='user')
    memory = relationship('MemoryVector', back_populates='user')
    logs = relationship('Logs', back_populates='user')


class ChatMessage(Base):
    __tablename__ = 'chat_messages'

    id = Column(Integer, primary_key=True)
    user_id = Column(BigInteger, ForeignKey('users.telegram_id'))
    role = Column(String(20))  # 'user' or 'assistant'
    content = Column(Text)
    input_tokens = Column(Integer, nullable=True)
    output_tokens = Column(Integer, nullable=True)
    timestamp = Column(TIMESTAMP(timezone=True), server_default=func.now())

    user = relationship('User', back_populates='messages')


class Payment(Base):
    __tablename__ = 'payments'

    id = Column(Integer, primary_key=True)

    user_id = Column(BigInteger, ForeignKey('users.telegram_id'))
    amount_usd = Column(Integer, nullable=False)
    crypto_amount = Column(String)
    crypto_currency = Column(String(20))  # 'TON', 'EVI'
    random_suffix = Column(String(10), nullable=False)
    status = Column(String(20), default='pending')  # pending, confirmed, failed
    created_at = Column(TIMESTAMP(timezone=True), server_default=func.now())
    confirmed_at = Column(TIMESTAMP(timezone=True), nullable=True)

    user = relationship('User', back_populates='payments')


class TokenPrice(Base):
    __tablename__ = 'token_prices'

    id = Column(Integer, primary_key=True)
    token = Column(String(20), unique=True)
    price_usd = Column(Float, nullable=False)
    updated_at = Column(TIMESTAMP(timezone=True), server_default=func.now(), server_onupdate=func.now())


class Wallet(Base):
    __tablename__ = 'wallets'

    id = Column(Integer, primary_key=True)
    user_id = Column(BigInteger, ForeignKey('users.telegram_id'))
    encrypted_private_key = Column(Text, nullable=False)
    created_at = Column(TIMESTAMP(timezone=True), server_default=func.now())

    user = relationship('User', back_populates='wallets')


class KnowledgeVector(Base):
    __tablename__ = 'knowledge_vectors'

    id = Column(Integer, primary_key=True)
    id_vector = Column(Text, nullable=False)
    uploaded_at = Column(TIMESTAMP(timezone=True), server_default=func.now())


class MemoryVector(Base):
    __tablename__ = 'memory_vectors'

    id = Column(Integer, primary_key=True)
    user_id = Column(BigInteger, ForeignKey('users.telegram_id'))
    id_vector = Column(Text, nullable=False)
    created_at = Column(TIMESTAMP(timezone=True), server_default=func.now())

    user = relationship('User', back_populates='memory')


class Logs(Base):
    __tablename__ = 'logs'

    id = Column(Integer, primary_key=True)
    user_id = Column(BigInteger, ForeignKey('users.telegram_id'))
    message = Column(Text, nullable=False)
    created_at = Column(TIMESTAMP(timezone=True), server_default=func.now())

    user = relationship('User', back_populates='logs')


async def create_tables():
    async with engine.begin() as conn:
        # await conn.run_sync(Base.metadata.drop_all)
        await conn.run_sync(Base.metadata.create_all)