files.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126
  1. from pydantic import BaseModel, ConfigDict
  2. from typing import Union, Optional
  3. import time
  4. import logging
  5. from sqlalchemy import Column, String, BigInteger, Text
  6. from apps.webui.internal.db import JSONField, Base, get_db
  7. import json
  8. from config import SRC_LOG_LEVELS
  9. log = logging.getLogger(__name__)
  10. log.setLevel(SRC_LOG_LEVELS["MODELS"])
  11. ####################
  12. # Files DB Schema
  13. ####################
  14. class File(Base):
  15. __tablename__ = "file"
  16. id = Column(String, primary_key=True)
  17. user_id = Column(String)
  18. filename = Column(Text)
  19. meta = Column(JSONField)
  20. created_at = Column(BigInteger)
  21. class FileModel(BaseModel):
  22. id: str
  23. user_id: str
  24. filename: str
  25. meta: dict
  26. created_at: int # timestamp in epoch
  27. model_config = ConfigDict(from_attributes=True)
  28. ####################
  29. # Forms
  30. ####################
  31. class FileModelResponse(BaseModel):
  32. id: str
  33. user_id: str
  34. filename: str
  35. meta: dict
  36. created_at: int # timestamp in epoch
  37. class FileForm(BaseModel):
  38. id: str
  39. filename: str
  40. meta: dict = {}
  41. class FilesTable:
  42. def insert_new_file(self, user_id: str, form_data: FileForm) -> Optional[FileModel]:
  43. with get_db() as db:
  44. file = FileModel(
  45. **{
  46. **form_data.model_dump(),
  47. "user_id": user_id,
  48. "created_at": int(time.time()),
  49. }
  50. )
  51. try:
  52. result = File(**file.model_dump())
  53. db.add(result)
  54. db.commit()
  55. db.refresh(result)
  56. if result:
  57. return FileModel.model_validate(result)
  58. else:
  59. return None
  60. except Exception as e:
  61. print(f"Error creating tool: {e}")
  62. return None
  63. def get_file_by_id(self, id: str) -> Optional[FileModel]:
  64. with get_db() as db:
  65. try:
  66. file = db.get(File, id)
  67. return FileModel.model_validate(file)
  68. except Exception:
  69. return None
  70. def get_files(self) -> list[FileModel]:
  71. with get_db() as db:
  72. return [FileModel.model_validate(file) for file in db.query(File).all()]
  73. def delete_file_by_id(self, id: str) -> bool:
  74. with get_db() as db:
  75. try:
  76. db.query(File).filter_by(id=id).delete()
  77. db.commit()
  78. return True
  79. except Exception:
  80. return False
  81. def delete_all_files(self) -> bool:
  82. with get_db() as db:
  83. try:
  84. db.query(File).delete()
  85. db.commit()
  86. return True
  87. except Exception:
  88. return False
  89. Files = FilesTable()