notes.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. import json
  2. import time
  3. import uuid
  4. from typing import Optional
  5. from open_webui.internal.db import Base, get_db
  6. from open_webui.utils.access_control import has_access
  7. from open_webui.models.users import Users, UserResponse
  8. from pydantic import BaseModel, ConfigDict
  9. from sqlalchemy import BigInteger, Boolean, Column, String, Text, JSON
  10. from sqlalchemy import or_, func, select, and_, text
  11. from sqlalchemy.sql import exists
  12. ####################
  13. # Note DB Schema
  14. ####################
  15. class Note(Base):
  16. __tablename__ = "note"
  17. id = Column(Text, primary_key=True)
  18. user_id = Column(Text)
  19. title = Column(Text)
  20. data = Column(JSON, nullable=True)
  21. meta = Column(JSON, nullable=True)
  22. access_control = Column(JSON, nullable=True)
  23. created_at = Column(BigInteger)
  24. updated_at = Column(BigInteger)
  25. class NoteModel(BaseModel):
  26. model_config = ConfigDict(from_attributes=True)
  27. id: str
  28. user_id: str
  29. title: str
  30. data: Optional[dict] = None
  31. meta: Optional[dict] = None
  32. access_control: Optional[dict] = None
  33. created_at: int # timestamp in epoch
  34. updated_at: int # timestamp in epoch
  35. ####################
  36. # Forms
  37. ####################
  38. class NoteForm(BaseModel):
  39. title: str
  40. data: Optional[dict] = None
  41. meta: Optional[dict] = None
  42. access_control: Optional[dict] = None
  43. class NoteUserResponse(NoteModel):
  44. user: Optional[UserResponse] = None
  45. class NoteTable:
  46. def insert_new_note(
  47. self,
  48. form_data: NoteForm,
  49. user_id: str,
  50. ) -> Optional[NoteModel]:
  51. with get_db() as db:
  52. note = NoteModel(
  53. **{
  54. "id": str(uuid.uuid4()),
  55. "user_id": user_id,
  56. **form_data.model_dump(),
  57. "created_at": int(time.time_ns()),
  58. "updated_at": int(time.time_ns()),
  59. }
  60. )
  61. new_note = Note(**note.model_dump())
  62. db.add(new_note)
  63. db.commit()
  64. return note
  65. def get_notes(self) -> list[NoteModel]:
  66. with get_db() as db:
  67. notes = db.query(Note).order_by(Note.updated_at.desc()).all()
  68. return [NoteModel.model_validate(note) for note in notes]
  69. def get_notes_by_user_id(
  70. self, user_id: str, permission: str = "write"
  71. ) -> list[NoteModel]:
  72. notes = self.get_notes()
  73. return [
  74. note
  75. for note in notes
  76. if note.user_id == user_id
  77. or has_access(user_id, permission, note.access_control)
  78. ]
  79. def get_note_by_id(self, id: str) -> Optional[NoteModel]:
  80. with get_db() as db:
  81. note = db.query(Note).filter(Note.id == id).first()
  82. return NoteModel.model_validate(note) if note else None
  83. def update_note_by_id(self, id: str, form_data: NoteForm) -> Optional[NoteModel]:
  84. with get_db() as db:
  85. note = db.query(Note).filter(Note.id == id).first()
  86. if not note:
  87. return None
  88. note.title = form_data.title
  89. note.data = form_data.data
  90. note.meta = form_data.meta
  91. note.access_control = form_data.access_control
  92. note.updated_at = int(time.time_ns())
  93. db.commit()
  94. return NoteModel.model_validate(note) if note else None
  95. def delete_note_by_id(self, id: str):
  96. with get_db() as db:
  97. db.query(Note).filter(Note.id == id).delete()
  98. db.commit()
  99. return True
  100. Notes = NoteTable()