Skip to content

Commit c748c3e

Browse files
tjbcktjbck
committed
refac
1 parent 3330802 commit c748c3e

File tree

3 files changed

+107
-57
lines changed

3 files changed

+107
-57
lines changed

‎backend/open_webui/models/chats.py‎

Lines changed: 54 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -431,22 +431,31 @@ def update_chat_title_by_id(self, id: str, title: str) -> Optional[ChatModel]:
431431
def update_chat_tags_by_id(
432432
self, id: str, tags: list[str], user
433433
) -> Optional[ChatModel]:
434-
chat = self.get_chat_by_id(id)
435-
if chat is None:
436-
return None
434+
with get_db_context() as db:
435+
chat = db.get(Chat, id)
436+
if chat is None:
437+
return None
437438

438-
self.delete_all_tags_by_id_and_user_id(id, user.id)
439+
old_tags = chat.meta.get("tags", [])
440+
new_tags = [t for t in tags if t.replace(" ", "_").lower() != "none"]
441+
new_tag_ids = [t.replace(" ", "_").lower() for t in new_tags]
439442

440-
for tag in chat.meta.get("tags", []):
441-
if self.count_chats_by_tag_name_and_user_id(tag, user.id) == 0:
442-
Tags.delete_tag_by_name_and_user_id(tag, user.id)
443+
# Single meta update
444+
chat.meta = {**chat.meta, "tags": new_tag_ids}
445+
db.commit()
446+
db.refresh(chat)
447+
448+
# Batch-create any missing tag rows
449+
Tags.ensure_tags_exist(new_tags, user.id, db=db)
443450

444-
for tag_name in tags:
445-
if tag_name.lower() == "none":
446-
continue
451+
# Clean up orphaned old tags in one query
452+
removed = set(old_tags) - set(new_tag_ids)
453+
if removed:
454+
self.delete_orphan_tags_for_user(
455+
list(removed), user.id, db=db
456+
)
447457

448-
self.add_chat_tag_by_id_and_user_id_and_tag_name(id, user.id, tag_name)
449-
return self.get_chat_by_id(id)
458+
return ChatModel.model_validate(chat)
450459

451460
def get_chat_title_by_id(self, id: str) -> Optional[str]:
452461
chat = self.get_chat_by_id(id)
@@ -1267,8 +1276,8 @@ def get_chat_tags_by_id_and_user_id(
12671276
) -> list[TagModel]:
12681277
with get_db_context(db) as db:
12691278
chat = db.get(Chat, id)
1270-
tags = chat.meta.get("tags", [])
1271-
return [Tags.get_tag_by_name_and_user_id(tag, user_id) for tag in tags]
1279+
tag_ids = chat.meta.get("tags", [])
1280+
return Tags.get_tags_by_ids_and_user_id(tag_ids, user_id, db=db)
12721281

12731282
def get_chat_list_by_user_id_and_tag_name(
12741283
self,
@@ -1309,20 +1318,16 @@ def get_chat_list_by_user_id_and_tag_name(
13091318
def add_chat_tag_by_id_and_user_id_and_tag_name(
13101319
self, id: str, user_id: str, tag_name: str, db: Optional[Session] = None
13111320
) -> Optional[ChatModel]:
1312-
tag = Tags.get_tag_by_name_and_user_id(tag_name, user_id)
1313-
if tag is None:
1314-
tag = Tags.insert_new_tag(tag_name, user_id)
1321+
tag_id = tag_name.replace(" ", "_").lower()
1322+
Tags.ensure_tags_exist([tag_name], user_id, db=db)
13151323
try:
13161324
with get_db_context(db) as db:
13171325
chat = db.get(Chat, id)
1318-
1319-
tag_id = tag.id
13201326
if tag_id not in chat.meta.get("tags", []):
13211327
chat.meta = {
13221328
**chat.meta,
13231329
"tags": list(set(chat.meta.get("tags", []) + [tag_id])),
13241330
}
1325-
13261331
db.commit()
13271332
db.refresh(chat)
13281333
return ChatModel.model_validate(chat)
@@ -1332,40 +1337,55 @@ def add_chat_tag_by_id_and_user_id_and_tag_name(
13321337
def count_chats_by_tag_name_and_user_id(
13331338
self, tag_name: str, user_id: str, db: Optional[Session] = None
13341339
) -> int:
1335-
with get_db_context(db) as db: # Assuming `get_db()` returns a session object
1340+
with get_db_context(db) as db:
13361341
query = db.query(Chat).filter_by(user_id=user_id, archived=False)
1337-
1338-
# Normalize the tag_name for consistency
13391342
tag_id = tag_name.replace(" ", "_").lower()
13401343

13411344
if db.bind.dialect.name == "sqlite":
1342-
# SQLite JSON1 support for querying the tags inside the `meta` JSON field
13431345
query = query.filter(
13441346
text(
1345-
f"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)"
1347+
"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)"
13461348
)
13471349
).params(tag_id=tag_id)
1348-
13491350
elif db.bind.dialect.name == "postgresql":
1350-
# PostgreSQL JSONB support for querying the tags inside the `meta` JSON field
13511351
query = query.filter(
13521352
text(
13531353
"EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)"
13541354
)
13551355
).params(tag_id=tag_id)
1356-
13571356
else:
13581357
raise NotImplementedError(
13591358
f"Unsupported dialect: {db.bind.dialect.name}"
13601359
)
13611360

1362-
# Get the count of matching records
1363-
count = query.count()
1364-
1365-
# Debugging output for inspection
1366-
log.info(f"Count of chats for tag '{tag_name}': {count}")
1361+
return query.count()
13671362

1368-
return count
1363+
def delete_orphan_tags_for_user(
1364+
self,
1365+
tag_ids: list[str],
1366+
user_id: str,
1367+
threshold: int = 0,
1368+
db: Optional[Session] = None,
1369+
) -> None:
1370+
"""Delete tag rows from *tag_ids* that appear in at most *threshold*
1371+
non-archived chats for *user_id*. One query to find orphans, one to
1372+
delete them.
1373+
1374+
Use threshold=0 after a tag is already removed from a chat's meta.
1375+
Use threshold=1 when the chat itself is about to be deleted (the
1376+
referencing chat still exists at query time).
1377+
"""
1378+
if not tag_ids:
1379+
return
1380+
with get_db_context(db) as db:
1381+
orphans = []
1382+
for tag_id in tag_ids:
1383+
count = self.count_chats_by_tag_name_and_user_id(
1384+
tag_id, user_id, db=db
1385+
)
1386+
if count <= threshold:
1387+
orphans.append(tag_id)
1388+
Tags.delete_tags_by_ids_and_user_id(orphans, user_id, db=db)
13691389

13701390
def count_chats_by_folder_id_and_user_id(
13711391
self, folder_id: str, user_id: str, db: Optional[Session] = None

‎backend/open_webui/models/tags.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,5 +115,45 @@ def delete_tag_by_name_and_user_id(
115115
log.error(f"delete_tag: {e}")
116116
return False
117117

118+
def delete_tags_by_ids_and_user_id(
119+
self, ids: list[str], user_id: str, db: Optional[Session] = None
120+
) -> bool:
121+
"""Delete all tags whose id is in *ids* for the given user, in one query."""
122+
if not ids:
123+
return True
124+
try:
125+
with get_db_context(db) as db:
126+
db.query(Tag).filter(
127+
Tag.id.in_(ids), Tag.user_id == user_id
128+
).delete(synchronize_session=False)
129+
db.commit()
130+
return True
131+
except Exception as e:
132+
log.error(f"delete_tags_by_ids: {e}")
133+
return False
134+
135+
def ensure_tags_exist(
136+
self, names: list[str], user_id: str, db: Optional[Session] = None
137+
) -> None:
138+
"""Create tag rows for any *names* that don't already exist for *user_id*."""
139+
if not names:
140+
return
141+
ids = [n.replace(" ", "_").lower() for n in names]
142+
with get_db_context(db) as db:
143+
existing = {
144+
t.id
145+
for t in db.query(Tag.id)
146+
.filter(Tag.id.in_(ids), Tag.user_id == user_id)
147+
.all()
148+
}
149+
new_tags = [
150+
Tag(id=tag_id, name=name, user_id=user_id)
151+
for tag_id, name in zip(ids, names)
152+
if tag_id not in existing
153+
]
154+
if new_tags:
155+
db.add_all(new_tags)
156+
db.commit()
157+
118158

119159
Tags = TagTable()

‎backend/open_webui/routers/chats.py‎

Lines changed: 13 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1131,9 +1131,9 @@ async def delete_chat_by_id(
11311131
status_code=status.HTTP_404_NOT_FOUND,
11321132
detail=ERROR_MESSAGES.NOT_FOUND,
11331133
)
1134-
for tag in chat.meta.get("tags", []):
1135-
if Chats.count_chats_by_tag_name_and_user_id(tag, user.id, db=db) == 1:
1136-
Tags.delete_tag_by_name_and_user_id(tag, user.id, db=db)
1134+
Chats.delete_orphan_tags_for_user(
1135+
chat.meta.get("tags", []), user.id, threshold=1, db=db
1136+
)
11371137

11381138
result = Chats.delete_chat_by_id(id, db=db)
11391139

@@ -1153,9 +1153,9 @@ async def delete_chat_by_id(
11531153
status_code=status.HTTP_404_NOT_FOUND,
11541154
detail=ERROR_MESSAGES.NOT_FOUND,
11551155
)
1156-
for tag in chat.meta.get("tags", []):
1157-
if Chats.count_chats_by_tag_name_and_user_id(tag, user.id, db=db) == 1:
1158-
Tags.delete_tag_by_name_and_user_id(tag, user.id, db=db)
1156+
Chats.delete_orphan_tags_for_user(
1157+
chat.meta.get("tags", []), user.id, threshold=1, db=db
1158+
)
11591159

11601160
result = Chats.delete_chat_by_id_and_user_id(id, user.id, db=db)
11611161
return result
@@ -1317,21 +1317,13 @@ async def archive_chat_by_id(
13171317
if chat:
13181318
chat = Chats.toggle_chat_archive_by_id(id, db=db)
13191319

1320-
# Delete tags if chat is archived
1320+
tag_ids = chat.meta.get("tags", [])
13211321
if chat.archived:
1322-
for tag_id in chat.meta.get("tags", []):
1323-
if (
1324-
Chats.count_chats_by_tag_name_and_user_id(tag_id, user.id, db=db)
1325-
== 0
1326-
):
1327-
log.debug(f"deleting tag: {tag_id}")
1328-
Tags.delete_tag_by_name_and_user_id(tag_id, user.id, db=db)
1322+
# Archived chats are excluded from count — clean up orphans
1323+
Chats.delete_orphan_tags_for_user(tag_ids, user.id, db=db)
13291324
else:
1330-
for tag_id in chat.meta.get("tags", []):
1331-
tag = Tags.get_tag_by_name_and_user_id(tag_id, user.id, db=db)
1332-
if tag is None:
1333-
log.debug(f"inserting tag: {tag_id}")
1334-
tag = Tags.insert_new_tag(tag_id, user.id, db=db)
1325+
# Unarchived — ensure tag rows exist
1326+
Tags.ensure_tags_exist(tag_ids, user.id, db=db)
13351327

13361328
return ChatResponse(**chat.model_dump())
13371329
else:
@@ -1537,11 +1529,9 @@ async def delete_all_tags_by_id(
15371529
):
15381530
chat = Chats.get_chat_by_id_and_user_id(id, user.id, db=db)
15391531
if chat:
1532+
old_tags = chat.meta.get("tags", [])
15401533
Chats.delete_all_tags_by_id_and_user_id(id, user.id, db=db)
1541-
1542-
for tag in chat.meta.get("tags", []):
1543-
if Chats.count_chats_by_tag_name_and_user_id(tag, user.id, db=db) == 0:
1544-
Tags.delete_tag_by_name_and_user_id(tag, user.id, db=db)
1534+
Chats.delete_orphan_tags_for_user(old_tags, user.id, db=db)
15451535

15461536
return True
15471537
else:

0 commit comments

Comments
 (0)