projects.py 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. import json
  2. import logging
  3. import time
  4. from typing import Optional
  5. from open_webui.apps.webui.internal.db import Base, get_db
  6. from open_webui.env import SRC_LOG_LEVELS
  7. from pydantic import BaseModel, ConfigDict
  8. from sqlalchemy import BigInteger, Column, String, Text, JSON
  9. log = logging.getLogger(__name__)
  10. log.setLevel(SRC_LOG_LEVELS["MODELS"])
  11. ####################
  12. # Projects DB Schema
  13. ####################
  14. class Project(Base):
  15. __tablename__ = "project"
  16. id = Column(Text, unique=True, primary_key=True)
  17. user_id = Column(Text)
  18. name = Column(Text)
  19. description = Column(Text)
  20. data = Column(JSON, nullable=True)
  21. meta = Column(JSON, nullable=True)
  22. created_at = Column(BigInteger)
  23. updated_at = Column(BigInteger)
  24. class ProjectModel(BaseModel):
  25. model_config = ConfigDict(from_attributes=True)
  26. id: str
  27. user_id: str
  28. name: str
  29. description: str
  30. data: Optional[dict] = None
  31. meta: Optional[dict] = None
  32. created_at: int # timestamp in epoch
  33. updated_at: int # timestamp in epoch
  34. ####################
  35. # Forms
  36. ####################
  37. class ProjectResponse(BaseModel):
  38. id: str
  39. name: str
  40. description: str
  41. data: Optional[dict] = None
  42. meta: Optional[dict] = None
  43. created_at: int # timestamp in epoch
  44. updated_at: int # timestamp in epoch
  45. class ProjectForm(BaseModel):
  46. id: str
  47. name: str
  48. description: str
  49. data: Optional[dict] = None
  50. class ProjectTable:
  51. def insert_new_project(
  52. self, user_id: str, form_data: ProjectForm
  53. ) -> Optional[ProjectModel]:
  54. with get_db() as db:
  55. project = ProjectModel(
  56. **{
  57. **form_data.model_dump(),
  58. "user_id": user_id,
  59. "created_at": int(time.time()),
  60. "updated_at": int(time.time()),
  61. }
  62. )
  63. try:
  64. result = Project(**project.model_dump())
  65. db.add(result)
  66. db.commit()
  67. db.refresh(result)
  68. if result:
  69. return ProjectModel.model_validate(result)
  70. else:
  71. return None
  72. except Exception:
  73. return None
  74. def get_projects(self) -> list[ProjectModel]:
  75. with get_db() as db:
  76. return [
  77. ProjectModel.model_validate(project)
  78. for project in db.query(Project).all()
  79. ]
  80. def get_project_by_id(self, id: str) -> Optional[ProjectModel]:
  81. try:
  82. with get_db() as db:
  83. project = db.query(Project).filter_by(id=id).first()
  84. return ProjectModel.model_validate(project) if project else None
  85. except Exception:
  86. return None
  87. def update_project_by_id(
  88. self, id: str, form_data: ProjectForm
  89. ) -> Optional[ProjectModel]:
  90. try:
  91. with get_db() as db:
  92. db.query(Project).filter_by(id=id).update(
  93. {
  94. "name": form_data.name,
  95. "updated_id": int(time.time()),
  96. }
  97. )
  98. db.commit()
  99. return self.get_project_by_id(id=form_data.id)
  100. except Exception as e:
  101. log.exception(e)
  102. return None
  103. def delete_project_by_id(self, id: str) -> bool:
  104. try:
  105. with get_db() as db:
  106. db.query(Project).filter_by(id=id).delete()
  107. db.commit()
  108. return True
  109. except Exception:
  110. return False
  111. Projects = ProjectTable()