diff --git a/backend/open_webui/apps/webui/models/documents.py b/backend/open_webui/apps/webui/models/documents.py deleted file mode 100644 index 0b96c2574..000000000 --- a/backend/open_webui/apps/webui/models/documents.py +++ /dev/null @@ -1,157 +0,0 @@ -import json -import logging -import time -from typing import Optional - -from open_webui.apps.webui.internal.db import Base, get_db -from open_webui.env import SRC_LOG_LEVELS -from pydantic import BaseModel, ConfigDict -from sqlalchemy import BigInteger, Column, String, Text - -log = logging.getLogger(__name__) -log.setLevel(SRC_LOG_LEVELS["MODELS"]) - -#################### -# Documents DB Schema -#################### - - -class Document(Base): - __tablename__ = "document" - - collection_name = Column(String, primary_key=True) - name = Column(String, unique=True) - title = Column(Text) - filename = Column(Text) - content = Column(Text, nullable=True) - user_id = Column(String) - timestamp = Column(BigInteger) - - -class DocumentModel(BaseModel): - model_config = ConfigDict(from_attributes=True) - - collection_name: str - name: str - title: str - filename: str - content: Optional[str] = None - user_id: str - timestamp: int # timestamp in epoch - - -#################### -# Forms -#################### - - -class DocumentResponse(BaseModel): - collection_name: str - name: str - title: str - filename: str - content: Optional[dict] = None - user_id: str - timestamp: int # timestamp in epoch - - -class DocumentUpdateForm(BaseModel): - name: str - title: str - - -class DocumentForm(DocumentUpdateForm): - collection_name: str - filename: str - content: Optional[str] = None - - -class DocumentsTable: - def insert_new_doc( - self, user_id: str, form_data: DocumentForm - ) -> Optional[DocumentModel]: - with get_db() as db: - document = DocumentModel( - **{ - **form_data.model_dump(), - "user_id": user_id, - "timestamp": int(time.time()), - } - ) - - try: - result = Document(**document.model_dump()) - db.add(result) - db.commit() - db.refresh(result) - if result: - return DocumentModel.model_validate(result) - else: - return None - except Exception: - return None - - def get_doc_by_name(self, name: str) -> Optional[DocumentModel]: - try: - with get_db() as db: - document = db.query(Document).filter_by(name=name).first() - return DocumentModel.model_validate(document) if document else None - except Exception: - return None - - def get_docs(self) -> list[DocumentModel]: - with get_db() as db: - return [ - DocumentModel.model_validate(doc) for doc in db.query(Document).all() - ] - - def update_doc_by_name( - self, name: str, form_data: DocumentUpdateForm - ) -> Optional[DocumentModel]: - try: - with get_db() as db: - db.query(Document).filter_by(name=name).update( - { - "title": form_data.title, - "name": form_data.name, - "timestamp": int(time.time()), - } - ) - db.commit() - return self.get_doc_by_name(form_data.name) - except Exception as e: - log.exception(e) - return None - - def update_doc_content_by_name( - self, name: str, updated: dict - ) -> Optional[DocumentModel]: - try: - doc = self.get_doc_by_name(name) - doc_content = json.loads(doc.content if doc.content else "{}") - doc_content = {**doc_content, **updated} - - with get_db() as db: - db.query(Document).filter_by(name=name).update( - { - "content": json.dumps(doc_content), - "timestamp": int(time.time()), - } - ) - db.commit() - return self.get_doc_by_name(name) - except Exception as e: - log.exception(e) - return None - - def delete_doc_by_name(self, name: str) -> bool: - try: - with get_db() as db: - db.query(Document).filter_by(name=name).delete() - db.commit() - return True - except Exception: - return False - - -Documents = DocumentsTable()