open-webui/backend/apps/webui/models/documents.py

from pydantic import BaseModel, ConfigDict
from typing import List, Optional
import time
import logging

from sqlalchemy import String, Column, BigInteger
from sqlalchemy.orm import Session

from apps.webui.internal.db import Base

import json

from config import SRC_LOG_LEVELS

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(String)
    filename = Column(String)
    content = Column(String, 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, db: Session, user_id: str, form_data: DocumentForm
    ) -> Optional[DocumentModel]:
        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:
            return None

    def get_doc_by_name(self, db: Session, name: str) -> Optional[DocumentModel]:
        try:
            document = db.query(Document).filter_by(name=name).first()
            return DocumentModel.model_validate(document) if document else None
        except:
            return None

    def get_docs(self, db: Session) -> List[DocumentModel]:
        return [DocumentModel.model_validate(doc) for doc in db.query(Document).all()]

    def update_doc_by_name(
        self, db: Session, name: str, form_data: DocumentUpdateForm
    ) -> Optional[DocumentModel]:
        try:
            db.query(Document).filter_by(name=name).update(
                {
                    "title": form_data.title,
                    "name": form_data.name,
                    "timestamp": int(time.time()),
                }
            )
            return self.get_doc_by_name(db, form_data.name)
        except Exception as e:
            log.exception(e)
            return None

    def update_doc_content_by_name(
        self, db: Session, name: str, updated: dict
    ) -> Optional[DocumentModel]:
        try:
            doc = self.get_doc_by_name(db, name)
            doc_content = json.loads(doc.content if doc.content else "{}")
            doc_content = {**doc_content, **updated}

            db.query(Document).filter_by(name=name).update(
                {
                    "content": json.dumps(doc_content),
                    "timestamp": int(time.time()),
                }
            )

            return self.get_doc_by_name(db, name)
        except Exception as e:
            log.exception(e)
            return None

    def delete_doc_by_name(self, db: Session, name: str) -> bool:
        try:
            db.query(Document).filter_by(name=name).delete()
            return True
        except:
            return False


Documents = DocumentsTable()