Mirrors/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2025-12-16 11:57:51 +01:00
Code Issues Packages Projects Releases Wiki Activity

feat(sqlalchemy): use session factory instead of context manager

Browse Source
This commit is contained in:
Jonathan Rohde
2024-06-24 13:06:15 +02:00
parent eb01e8d275
commit da403f3e3c
15 changed files with 640 additions and 759 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
backend/apps/webui/models/tools.py
View File

@@ -3,9 +3,8 @@ 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, JSONField, get_session
from apps.webui.internal.db import Base, JSONField, Session
from apps.webui.models.users import Users
import json
@@ -95,48 +94,43 @@ class ToolsTable:
)
try:
with get_session() as db:
result = Tool(**tool.model_dump())
db.add(result)
db.commit()
db.refresh(result)
if result:
return ToolModel.model_validate(result)
else:
return None
result = Tool(**tool.model_dump())
Session.add(result)
Session.commit()
Session.refresh(result)
if result:
return ToolModel.model_validate(result)
else:
return None
except Exception as e:
print(f"Error creating tool: {e}")
return None
def get_tool_by_id(self, id: str) -> Optional[ToolModel]:
try:
with get_session() as db:
tool = db.get(Tool, id)
return ToolModel.model_validate(tool)
tool = Session.get(Tool, id)
return ToolModel.model_validate(tool)
except:
return None
def get_tools(self) -> List[ToolModel]:
with get_session() as db:
return [ToolModel.model_validate(tool) for tool in db.query(Tool).all()]
return [ToolModel.model_validate(tool) for tool in Session.query(Tool).all()]
def get_tool_valves_by_id(self, id: str) -> Optional[dict]:
try:
with get_session() as db:
tool = db.get(Tool, id)
return tool.valves if tool.valves else {}
tool = Session.get(Tool, id)
return tool.valves if tool.valves else {}
except Exception as e:
print(f"An error occurred: {e}")
return None
def update_tool_valves_by_id(self, id: str, valves: dict) -> Optional[ToolValves]:
try:
with get_session() as db:
db.query(Tool).filter_by(id=id).update(
{"valves": valves, "updated_at": int(time.time())}
)
db.commit()
return self.get_tool_by_id(id)
Session.query(Tool).filter_by(id=id).update(
{"valves": valves, "updated_at": int(time.time())}
)
Session.commit()
return self.get_tool_by_id(id)
except:
return None
@@ -183,19 +177,18 @@ class ToolsTable:
def update_tool_by_id(self, id: str, updated: dict) -> Optional[ToolModel]:
try:
with get_session() as db:
db.query(Tool).filter_by(id=id).update(
{**updated, "updated_at": int(time.time())}
)
db.commit()
return self.get_tool_by_id(id)
tool = Session.get(Tool, id)
tool.update(**updated)
tool.updated_at = int(time.time())
Session.commit()
Session.refresh(tool)
return ToolModel.model_validate(tool)
except:
return None
def delete_tool_by_id(self, id: str) -> bool:
try:
with get_session() as db:
db.query(Tool).filter_by(id=id).delete()
Session.query(Tool).filter_by(id=id).delete()
return True
except:
return False