files.py 2.71 KB
Newer Older
1
from pydantic import BaseModel, ConfigDict
Timothy J. Baek's avatar
Timothy J. Baek committed
2
3
4
from typing import List, Union, Optional
import time
import logging
5
6
7
8

from sqlalchemy import Column, String, BigInteger
from sqlalchemy.orm import Session

9
from apps.webui.internal.db import JSONField, Base, get_session
Timothy J. Baek's avatar
Timothy J. Baek committed
10
11
12
13
14
15
16
17
18
19
20
21
22

import json

from config import SRC_LOG_LEVELS

log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])

####################
# Files DB Schema
####################


23
24
class File(Base):
    __tablename__ = "file"
Timothy J. Baek's avatar
Timothy J. Baek committed
25

26
27
28
29
30
    id = Column(String, primary_key=True)
    user_id = Column(String)
    filename = Column(String)
    meta = Column(JSONField)
    created_at = Column(BigInteger)
Timothy J. Baek's avatar
Timothy J. Baek committed
31
32
33
34
35
36
37
38
39


class FileModel(BaseModel):
    id: str
    user_id: str
    filename: str
    meta: dict
    created_at: int  # timestamp in epoch

40
    model_config = ConfigDict(from_attributes=True)
Timothy J. Baek's avatar
Timothy J. Baek committed
41

42

Timothy J. Baek's avatar
Timothy J. Baek committed
43
44
45
46
47
####################
# Forms
####################


Timothy J. Baek's avatar
Timothy J. Baek committed
48
class FileModelResponse(BaseModel):
Timothy J. Baek's avatar
Timothy J. Baek committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
    id: str
    user_id: str
    filename: str
    meta: dict
    created_at: int  # timestamp in epoch


class FileForm(BaseModel):
    id: str
    filename: str
    meta: dict = {}


class FilesTable:

64
    def insert_new_file(self, user_id: str, form_data: FileForm) -> Optional[FileModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
65
66
67
68
69
70
71
72
73
        file = FileModel(
            **{
                **form_data.model_dump(),
                "user_id": user_id,
                "created_at": int(time.time()),
            }
        )

        try:
74
75
76
77
78
79
80
81
82
            with get_session() as db:
                result = File(**file.model_dump())
                db.add(result)
                db.commit()
                db.refresh(result)
                if result:
                    return FileModel.model_validate(result)
                else:
                    return None
Timothy J. Baek's avatar
Timothy J. Baek committed
83
84
85
86
        except Exception as e:
            print(f"Error creating tool: {e}")
            return None

87
    def get_file_by_id(self, id: str) -> Optional[FileModel]:
Timothy J. Baek's avatar
Timothy J. Baek committed
88
        try:
89
90
91
            with get_session() as db:
                file = db.get(File, id)
                return FileModel.model_validate(file)
Timothy J. Baek's avatar
Timothy J. Baek committed
92
93
94
        except:
            return None

95
96
97
    def get_files(self) -> List[FileModel]:
        with get_session() as db:
            return [FileModel.model_validate(file) for file in db.query(File).all()]
Timothy J. Baek's avatar
Timothy J. Baek committed
98

99
    def delete_file_by_id(self, id: str) -> bool:
Timothy J. Baek's avatar
Timothy J. Baek committed
100
        try:
101
102
103
            with get_session() as db:
                db.query(File).filter_by(id=id).delete()
                db.commit()
Timothy J. Baek's avatar
Timothy J. Baek committed
104
105
106
107
            return True
        except:
            return False

108
    def delete_all_files(self) -> bool:
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
109
        try:
110
111
112
            with get_session() as db:
                db.query(File).delete()
                db.commit()
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
113
114
115
116
            return True
        except:
            return False

Timothy J. Baek's avatar
Timothy J. Baek committed
117

118
Files = FilesTable()