~cytrogen/evi-run

ref: 77f8b617bfe2a54b8478c461c43d6c0e6e6a07b5 evi-run/database/models.py -rw-r--r-- 4.1 KiB
77f8b617 — Bendy docs(readme): update GPT-5 examples and model_settings; replace hardcoded line refs 8 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)