projects.py 3.6 KB

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