open-webui/backend/apps/webui/models/files.py

127 lines
2.8 KiB
Python
Raw Normal View History

from pydantic import BaseModel, ConfigDict
2024-06-18 18:36:55 +00:00
from typing import List, Union, Optional
import time
import logging
from sqlalchemy import Column, String, BigInteger, Text
2024-07-04 06:32:39 +00:00
from apps.webui.internal.db import JSONField, Base, get_db
2024-06-18 18:36:55 +00:00
import json
from config import SRC_LOG_LEVELS
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MODELS"])
####################
# Files DB Schema
####################
class File(Base):
__tablename__ = "file"
2024-06-18 18:36:55 +00:00
id = Column(String, primary_key=True)
user_id = Column(String)
filename = Column(Text)
meta = Column(JSONField)
created_at = Column(BigInteger)
2024-06-18 18:36:55 +00:00
class FileModel(BaseModel):
id: str
user_id: str
filename: str
meta: dict
created_at: int # timestamp in epoch
model_config = ConfigDict(from_attributes=True)
2024-06-18 18:36:55 +00:00
2024-06-24 07:57:08 +00:00
2024-06-18 18:36:55 +00:00
####################
# Forms
####################
2024-06-18 21:33:44 +00:00
class FileModelResponse(BaseModel):
2024-06-18 18:36:55 +00:00
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:
def insert_new_file(self, user_id: str, form_data: FileForm) -> Optional[FileModel]:
2024-07-04 06:32:39 +00:00
with get_db() as db:
file = FileModel(
**{
**form_data.model_dump(),
"user_id": user_id,
"created_at": int(time.time()),
}
)
try:
result = File(**file.model_dump())
db.add(result)
db.commit()
db.refresh(result)
if result:
return FileModel.model_validate(result)
else:
return None
except Exception as e:
print(f"Error creating tool: {e}")
return None
2024-06-18 18:36:55 +00:00
def get_file_by_id(self, id: str) -> Optional[FileModel]:
2024-07-04 06:32:39 +00:00
with get_db() as db:
try:
file = db.get(File, id)
return FileModel.model_validate(file)
2024-08-14 12:38:19 +00:00
except Exception:
2024-07-04 06:32:39 +00:00
return None
2024-06-18 18:36:55 +00:00
def get_files(self) -> List[FileModel]:
2024-07-04 06:32:39 +00:00
with get_db() as db:
return [FileModel.model_validate(file) for file in db.query(File).all()]
2024-06-18 18:36:55 +00:00
def delete_file_by_id(self, id: str) -> bool:
2024-07-04 06:32:39 +00:00
with get_db() as db:
try:
db.query(File).filter_by(id=id).delete()
2024-07-06 15:10:58 +00:00
db.commit()
2024-07-04 06:32:39 +00:00
return True
2024-08-14 12:38:19 +00:00
except Exception:
2024-07-04 06:32:39 +00:00
return False
2024-06-18 18:36:55 +00:00
def delete_all_files(self) -> bool:
2024-07-04 06:32:39 +00:00
with get_db() as db:
try:
db.query(File).delete()
2024-07-06 15:10:58 +00:00
db.commit()
2024-07-04 06:32:39 +00:00
return True
2024-08-14 12:38:19 +00:00
except Exception:
2024-07-04 06:32:39 +00:00
return False
2024-06-18 21:15:08 +00:00
2024-06-18 18:36:55 +00:00
Files = FilesTable()