feat(sqlalchemy): format backend

open-webui · tjbck · Jul 6, 2024 · Jun 18, 2024 · Jun 21, 2024 · Jun 24, 2024
commit c134eab27a929cbf678a60356a4c8f6c2e718201
diff --git a/backend/apps/webui/internal/db.py b/backend/apps/webui/internal/db.py
@@ -53,7 53,9 @@ def python_value(self, value):
  )
 else:
  engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine, expire_on_commit=False)
 SessionLocal = sessionmaker(
  autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
 )
 Base = declarative_base()
 
 
@@ -66,4 68,3 @@ def get_session():
  except Exception as e:
  db.rollback()
  raise e
-
diff --git a/backend/apps/webui/models/auths.py b/backend/apps/webui/models/auths.py
@@ -126,9 126,7 @@ def insert_new_auth(
  else:
  return None
 
- def authenticate_user(
- self, email: str, password: str
- ) -> Optional[UserModel]:
  def authenticate_user(self, email: str, password: str) -> Optional[UserModel]:
  log.info(f"authenticate_user: {email}")
  with get_session() as db:
  try:
@@ -144,9 142,7 @@ def authenticate_user(
  except:
  return None
 
- def authenticate_user_by_api_key(
- self, api_key: str
- ) -> Optional[UserModel]:
  def authenticate_user_by_api_key(self, api_key: str) -> Optional[UserModel]:
  log.info(f"authenticate_user_by_api_key: {api_key}")
  with get_session() as db:
  # if no api_key, return None
@@ -159,9 155,7 @@ def authenticate_user_by_api_key(
  except:
  return False
 
- def authenticate_user_by_trusted_header(
- self, email: str
- ) -> Optional[UserModel]:
  def authenticate_user_by_trusted_header(self, email: str) -> Optional[UserModel]:
  log.info(f"authenticate_user_by_trusted_header: {email}")
  with get_session() as db:
  try:
@@ -172,12 166,12 @@ def authenticate_user_by_trusted_header(
  except:
  return None
 
- def update_user_password_by_id(
- self, id: str, new_password: str
- ) -> bool:
  def update_user_password_by_id(self, id: str, new_password: str) -> bool:
  with get_session() as db:
  try:
- result = db.query(Auth).filter_by(id=id).update({"password": new_password})
  result = (
  db.query(Auth).filter_by(id=id).update({"password": new_password})
  )
  return True if result == 1 else False
  except:
  return False

diff --git a/backend/apps/webui/models/chats.py b/backend/apps/webui/models/chats.py
@@ -79,17 79,17 @@ class ChatTitleIdResponse(BaseModel):
 
 class ChatTable:
 
- def insert_new_chat(
- self, user_id: str, form_data: ChatForm
- ) -> Optional[ChatModel]:
  def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]:
  with get_session() as db:
  id = str(uuid.uuid4())
  chat = ChatModel(
  **{
  "id": id,
  "user_id": user_id,
  "title": (
- form_data.chat["title"] if "title" in form_data.chat else "New Chat"
  form_data.chat["title"]
  if "title" in form_data.chat
  else "New Chat"
  ),
  "chat": json.dumps(form_data.chat),
  "created_at": int(time.time()),
@@ -103,9 103,7 @@ def insert_new_chat(
  db.refresh(result)
  return ChatModel.model_validate(result) if result else None
 
- def update_chat_by_id(
- self, id: str, chat: dict
- ) -> Optional[ChatModel]:
  def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]:
  with get_session() as db:
  try:
  chat_obj = db.get(Chat, id)
@@ -119,9 117,7 @@ def update_chat_by_id(
  except Exception as e:
  return None
 
- def insert_shared_chat_by_chat_id(
- self, chat_id: str
- ) -> Optional[ChatModel]:
  def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
  with get_session() as db:
  # Get the existing chat to share
  chat = db.get(Chat, chat_id)
@@ -145,14 141,14 @@ def insert_shared_chat_by_chat_id(
  db.refresh(shared_result)
  # Update the original chat with the share_id
  result = (
- db.query(Chat).filter_by(id=chat_id).update({"share_id": shared_chat.id})
  db.query(Chat)
  .filter_by(id=chat_id)
  .update({"share_id": shared_chat.id})
  )
 
  return shared_chat if (shared_result and result) else None
 
- def update_shared_chat_by_chat_id(
- self, chat_id: str
- ) -> Optional[ChatModel]:
  def update_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
  with get_session() as db:
  try:
  print("update_shared_chat_by_id")
@@ -271,9 267,7 @@ def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]:
  except Exception as e:
  return None
 
- def get_chat_by_id_and_user_id(
- self, id: str, user_id: str
- ) -> Optional[ChatModel]:
  def get_chat_by_id_and_user_id(self, id: str, user_id: str) -> Optional[ChatModel]:
  try:
  with get_session() as db:
  chat = db.query(Chat).filter_by(id=id, user_id=user_id).first()
@@ -293,13 287,13 @@ def get_chats(self, skip: int = 0, limit: int = 50) -> List[ChatModel]:
  def get_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
  with get_session() as db:
  all_chats = (
- db.query(Chat).filter_by(user_id=user_id).order_by(Chat.updated_at.desc())
  db.query(Chat)
  .filter_by(user_id=user_id)
  .order_by(Chat.updated_at.desc())
  )
  return [ChatModel.model_validate(chat) for chat in all_chats]
 
- def get_archived_chats_by_user_id(
- self, user_id: str
- ) -> List[ChatModel]:
  def get_archived_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
  with get_session() as db:
  all_chats = (
  db.query(Chat)

diff --git a/backend/apps/webui/models/documents.py b/backend/apps/webui/models/documents.py
@@ -106,7 106,9 @@ def get_doc_by_name(self, name: str) -> Optional[DocumentModel]:
 
  def get_docs(self) -> List[DocumentModel]:
  with get_session() as db:
- return [DocumentModel.model_validate(doc) for doc in db.query(Document).all()]
  return [
  DocumentModel.model_validate(doc) for doc in db.query(Document).all()
  ]
 
  def update_doc_by_name(
  self, name: str, form_data: DocumentUpdateForm

diff --git a/backend/apps/webui/models/files.py b/backend/apps/webui/models/files.py
@@ -39,6 39,7 @@ class FileModel(BaseModel):
 
  model_config = ConfigDict(from_attributes=True)
 
 
 ####################
 # Forms
 ####################

diff --git a/backend/apps/webui/models/functions.py b/backend/apps/webui/models/functions.py
@@ -142,9 142,9 @@ def get_functions_by_type(
  with get_session() as db:
  return [
  FunctionModel.model_validate(function)
- for function in db.query(Function).filter_by(
-  type=type, is_active=True
- ).all()
  for function in db.query(Function)
  .filter_by(type=type, is_active=True)
  .all()
  ]
  else:
  with get_session() as db:
@@ -220,10 220,12 @@ def update_user_valves_by_id_and_user_id(
  def update_function_by_id(self, id: str, updated: dict) -> Optional[FunctionModel]:
  try:
  with get_session() as db:
- db.query(Function).filter_by(id=id).update({
- **updated,
- "updated_at": int(time.time()),
- })
  db.query(Function).filter_by(id=id).update(
  {
  **updated,
  "updated_at": int(time.time()),
  }
  )
  db.commit()
  return self.get_function_by_id(id)
  except:
@@ -232,10 234,12 @@ def update_function_by_id(self, id: str, updated: dict) -> Optional[FunctionMode
  def deactivate_all_functions(self) -> Optional[bool]:
  try:
  with get_session() as db:
- db.query(Function).update({
- "is_active": False,
- "updated_at": int(time.time()),
- })
  db.query(Function).update(
  {
  "is_active": False,
  "updated_at": int(time.time()),
  }
  )
  db.commit()
  return True
  except:

diff --git a/backend/apps/webui/models/models.py b/backend/apps/webui/models/models.py
@@ -153,9 153,7 @@ def get_model_by_id(self, id: str) -> Optional[ModelModel]:
  except:
  return None
 
- def update_model_by_id(
- self, id: str, model: ModelForm
- ) -> Optional[ModelModel]:
  def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]:
  try:
  # update only the fields that are present in the model
  with get_session() as db:

diff --git a/backend/apps/webui/models/prompts.py b/backend/apps/webui/models/prompts.py
@@ -83,7 83,9 @@ def get_prompt_by_command(self, command: str) -> Optional[PromptModel]:
 
  def get_prompts(self) -> List[PromptModel]:
  with get_session() as db:
- return [PromptModel.model_validate(prompt) for prompt in db.query(Prompt).all()]
  return [
  PromptModel.model_validate(prompt) for prompt in db.query(Prompt).all()
  ]
 
  def update_prompt_by_command(
  self, command: str, form_data: PromptForm

diff --git a/backend/apps/webui/models/tags.py b/backend/apps/webui/models/tags.py
@@ -79,9 79,7 @@ class ChatTagsResponse(BaseModel):
 
 class TagTable:
 
- def insert_new_tag(
- self, name: str, user_id: str
- ) -> Optional[TagModel]:
  def insert_new_tag(self, name: str, user_id: str) -> Optional[TagModel]:
  id = str(uuid.uuid4())
  tag = TagModel(**{"id": id, "user_id": user_id, "name": name})
  try:
@@ -201,11 199,13 @@ def count_chat_ids_by_tag_name_and_user_id(
  self, tag_name: str, user_id: str
  ) -> int:
  with get_session() as db:
- return db.query(ChatIdTag).filter_by(tag_name=tag_name, user_id=user_id).count()
  return (
  db.query(ChatIdTag)
  .filter_by(tag_name=tag_name, user_id=user_id)
  .count()
  )
 
- def delete_tag_by_tag_name_and_user_id(
- self, tag_name: str, user_id: str
- ) -> bool:
  def delete_tag_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> bool:
  try:
  with get_session() as db:
  res = (
@@ -252,9 252,7 @@ def delete_tag_by_tag_name_and_chat_id_and_user_id(
  log.error(f"delete_tag: {e}")
  return False
 
- def delete_tags_by_chat_id_and_user_id(
- self, chat_id: str, user_id: str
- ) -> bool:
  def delete_tags_by_chat_id_and_user_id(self, chat_id: str, user_id: str) -> bool:
  tags = self.get_tags_by_chat_id_and_user_id(chat_id, user_id)
 
  for tag in tags:

diff --git a/backend/apps/webui/models/users.py b/backend/apps/webui/models/users.py
@@ -165,9 165,7 @@ def get_first_user(self) -> UserModel:
  except:
  return None
 
- def update_user_role_by_id(
- self, id: str, role: str
- ) -> Optional[UserModel]:
  def update_user_role_by_id(self, id: str, role: str) -> Optional[UserModel]:
  with get_session() as db:
  try:
  db.query(User).filter_by(id=id).update({"role": role})
@@ -193,12 191,12 @@ def update_user_profile_image_url_by_id(
  except:
  return None
 
- def update_user_last_active_by_id(
- self, id: str
- ) -> Optional[UserModel]:
  def update_user_last_active_by_id(self, id: str) -> Optional[UserModel]:
  with get_session() as db:
  try:
- db.query(User).filter_by(id=id).update({"last_active_at": int(time.time())})
  db.query(User).filter_by(id=id).update(
  {"last_active_at": int(time.time())}
  )
 
  user = db.query(User).filter_by(id=id).first()
  return UserModel.model_validate(user)
@@ -217,9 215,7 @@ def update_user_oauth_sub_by_id(
  except:
  return None
 
- def update_user_by_id(
- self, id: str, updated: dict
- ) -> Optional[UserModel]:
  def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]:
  with get_session() as db:
  try:
  db.query(User).filter_by(id=id).update(updated)

diff --git a/backend/apps/webui/routers/auths.py b/backend/apps/webui/routers/auths.py
@@ -78,8 78,7 @@ async def get_session_user(
 
 @router.post("/update/profile", response_model=UserResponse)
 async def update_profile(
- form_data: UpdateProfileForm,
- session_user=Depends(get_current_user)
  form_data: UpdateProfileForm, session_user=Depends(get_current_user)
 ):
  if session_user:
  user = Users.update_user_by_id(
@@ -101,8 100,7 @@ async def update_profile(
 
 @router.post("/update/password", response_model=bool)
 async def update_password(
- form_data: UpdatePasswordForm,
- session_user=Depends(get_current_user)
  form_data: UpdatePasswordForm, session_user=Depends(get_current_user)
 ):
  if WEBUI_AUTH_TRUSTED_EMAIL_HEADER:
  raise HTTPException(400, detail=ERROR_MESSAGES.ACTION_PROHIBITED)
@@ -269,9 267,7 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
 
 
 @router.post("/add", response_model=SigninResponse)
-async def add_user(
- form_data: AddUserForm, user=Depends(get_admin_user)
-):
 async def add_user(form_data: AddUserForm, user=Depends(get_admin_user)):
 
  if not validate_email_format(form_data.email.lower()):
  raise HTTPException(
@@ -316,9 312,7 @@ async def add_user(
 
 
 @router.get("/admin/details")
-async def get_admin_details(
- request: Request, user=Depends(get_current_user)
-):
 async def get_admin_details(request: Request, user=Depends(get_current_user)):
  if request.app.state.config.SHOW_ADMIN_DETAILS:
  admin_email = request.app.state.config.ADMIN_EMAIL
  admin_name = None