chats.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908
  1. import json
  2. import time
  3. import uuid
  4. from typing import Optional
  5. from open_webui.internal.db import Base, get_db
  6. from open_webui.models.tags import TagModel, Tag, Tags
  7. from pydantic import BaseModel, ConfigDict
  8. from sqlalchemy import BigInteger, Boolean, Column, String, Text, JSON
  9. from sqlalchemy import or_, func, select, and_, text
  10. from sqlalchemy.sql import exists
  11. ####################
  12. # Chat DB Schema
  13. ####################
  14. class Chat(Base):
  15. __tablename__ = "chat"
  16. id = Column(String, primary_key=True)
  17. user_id = Column(String)
  18. title = Column(Text)
  19. chat = Column(JSON)
  20. created_at = Column(BigInteger)
  21. updated_at = Column(BigInteger)
  22. share_id = Column(Text, unique=True, nullable=True)
  23. archived = Column(Boolean, default=False)
  24. pinned = Column(Boolean, default=False, nullable=True)
  25. meta = Column(JSON, server_default="{}")
  26. folder_id = Column(Text, nullable=True)
  27. class ChatModel(BaseModel):
  28. model_config = ConfigDict(from_attributes=True)
  29. id: str
  30. user_id: str
  31. title: str
  32. chat: dict
  33. created_at: int # timestamp in epoch
  34. updated_at: int # timestamp in epoch
  35. share_id: Optional[str] = None
  36. archived: bool = False
  37. pinned: Optional[bool] = False
  38. meta: dict = {}
  39. folder_id: Optional[str] = None
  40. ####################
  41. # Forms
  42. ####################
  43. class ChatForm(BaseModel):
  44. chat: dict
  45. class ChatImportForm(ChatForm):
  46. meta: Optional[dict] = {}
  47. pinned: Optional[bool] = False
  48. folder_id: Optional[str] = None
  49. class ChatTitleMessagesForm(BaseModel):
  50. title: str
  51. messages: list[dict]
  52. class ChatTitleForm(BaseModel):
  53. title: str
  54. class ChatResponse(BaseModel):
  55. id: str
  56. user_id: str
  57. title: str
  58. chat: dict
  59. updated_at: int # timestamp in epoch
  60. created_at: int # timestamp in epoch
  61. share_id: Optional[str] = None # id of the chat to be shared
  62. archived: bool
  63. pinned: Optional[bool] = False
  64. meta: dict = {}
  65. folder_id: Optional[str] = None
  66. class ChatTitleIdResponse(BaseModel):
  67. id: str
  68. title: str
  69. updated_at: int
  70. created_at: int
  71. class ChatTable:
  72. def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]:
  73. with get_db() as db:
  74. id = str(uuid.uuid4())
  75. chat = ChatModel(
  76. **{
  77. "id": id,
  78. "user_id": user_id,
  79. "title": (
  80. form_data.chat["title"]
  81. if "title" in form_data.chat
  82. else "New Chat"
  83. ),
  84. "chat": form_data.chat,
  85. "created_at": int(time.time()),
  86. "updated_at": int(time.time()),
  87. }
  88. )
  89. result = Chat(**chat.model_dump())
  90. db.add(result)
  91. db.commit()
  92. db.refresh(result)
  93. return ChatModel.model_validate(result) if result else None
  94. def import_chat(
  95. self, user_id: str, form_data: ChatImportForm
  96. ) -> Optional[ChatModel]:
  97. with get_db() as db:
  98. id = str(uuid.uuid4())
  99. chat = ChatModel(
  100. **{
  101. "id": id,
  102. "user_id": user_id,
  103. "title": (
  104. form_data.chat["title"]
  105. if "title" in form_data.chat
  106. else "New Chat"
  107. ),
  108. "chat": form_data.chat,
  109. "meta": form_data.meta,
  110. "pinned": form_data.pinned,
  111. "folder_id": form_data.folder_id,
  112. "created_at": int(time.time()),
  113. "updated_at": int(time.time()),
  114. }
  115. )
  116. result = Chat(**chat.model_dump())
  117. db.add(result)
  118. db.commit()
  119. db.refresh(result)
  120. return ChatModel.model_validate(result) if result else None
  121. def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]:
  122. try:
  123. with get_db() as db:
  124. chat_item = db.get(Chat, id)
  125. chat_item.chat = chat
  126. chat_item.title = chat["title"] if "title" in chat else "New Chat"
  127. chat_item.updated_at = int(time.time())
  128. db.commit()
  129. db.refresh(chat_item)
  130. return ChatModel.model_validate(chat_item)
  131. except Exception:
  132. return None
  133. def update_chat_title_by_id(self, id: str, title: str) -> Optional[ChatModel]:
  134. chat = self.get_chat_by_id(id)
  135. if chat is None:
  136. return None
  137. chat = chat.chat
  138. chat["title"] = title
  139. return self.update_chat_by_id(id, chat)
  140. def update_chat_tags_by_id(
  141. self, id: str, tags: list[str], user
  142. ) -> Optional[ChatModel]:
  143. chat = self.get_chat_by_id(id)
  144. if chat is None:
  145. return None
  146. self.delete_all_tags_by_id_and_user_id(id, user.id)
  147. for tag in chat.meta.get("tags", []):
  148. if self.count_chats_by_tag_name_and_user_id(tag, user.id) == 0:
  149. Tags.delete_tag_by_name_and_user_id(tag, user.id)
  150. for tag_name in tags:
  151. if tag_name.lower() == "none":
  152. continue
  153. self.add_chat_tag_by_id_and_user_id_and_tag_name(id, user.id, tag_name)
  154. return self.get_chat_by_id(id)
  155. def get_chat_title_by_id(self, id: str) -> Optional[str]:
  156. chat = self.get_chat_by_id(id)
  157. if chat is None:
  158. return None
  159. return chat.chat.get("title", "New Chat")
  160. def get_messages_by_chat_id(self, id: str) -> Optional[dict]:
  161. chat = self.get_chat_by_id(id)
  162. if chat is None:
  163. return None
  164. return chat.chat.get("history", {}).get("messages", {}) or {}
  165. def get_message_by_id_and_message_id(
  166. self, id: str, message_id: str
  167. ) -> Optional[dict]:
  168. chat = self.get_chat_by_id(id)
  169. if chat is None:
  170. return None
  171. return chat.chat.get("history", {}).get("messages", {}).get(message_id, {})
  172. def upsert_message_to_chat_by_id_and_message_id(
  173. self, id: str, message_id: str, message: dict
  174. ) -> Optional[ChatModel]:
  175. chat = self.get_chat_by_id(id)
  176. if chat is None:
  177. return None
  178. chat = chat.chat
  179. history = chat.get("history", {})
  180. if message_id in history.get("messages", {}):
  181. history["messages"][message_id] = {
  182. **history["messages"][message_id],
  183. **message,
  184. }
  185. else:
  186. history["messages"][message_id] = message
  187. history["currentId"] = message_id
  188. chat["history"] = history
  189. return self.update_chat_by_id(id, chat)
  190. def add_message_status_to_chat_by_id_and_message_id(
  191. self, id: str, message_id: str, status: dict
  192. ) -> Optional[ChatModel]:
  193. chat = self.get_chat_by_id(id)
  194. if chat is None:
  195. return None
  196. chat = chat.chat
  197. history = chat.get("history", {})
  198. if message_id in history.get("messages", {}):
  199. status_history = history["messages"][message_id].get("statusHistory", [])
  200. status_history.append(status)
  201. history["messages"][message_id]["statusHistory"] = status_history
  202. chat["history"] = history
  203. return self.update_chat_by_id(id, chat)
  204. def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
  205. with get_db() as db:
  206. # Get the existing chat to share
  207. chat = db.get(Chat, chat_id)
  208. # Check if the chat is already shared
  209. if chat.share_id:
  210. return self.get_chat_by_id_and_user_id(chat.share_id, "shared")
  211. # Create a new chat with the same data, but with a new ID
  212. shared_chat = ChatModel(
  213. **{
  214. "id": str(uuid.uuid4()),
  215. "user_id": f"shared-{chat_id}",
  216. "title": chat.title,
  217. "chat": chat.chat,
  218. "created_at": chat.created_at,
  219. "updated_at": int(time.time()),
  220. }
  221. )
  222. shared_result = Chat(**shared_chat.model_dump())
  223. db.add(shared_result)
  224. db.commit()
  225. db.refresh(shared_result)
  226. # Update the original chat with the share_id
  227. result = (
  228. db.query(Chat)
  229. .filter_by(id=chat_id)
  230. .update({"share_id": shared_chat.id})
  231. )
  232. db.commit()
  233. return shared_chat if (shared_result and result) else None
  234. def update_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
  235. try:
  236. with get_db() as db:
  237. chat = db.get(Chat, chat_id)
  238. shared_chat = (
  239. db.query(Chat).filter_by(user_id=f"shared-{chat_id}").first()
  240. )
  241. if shared_chat is None:
  242. return self.insert_shared_chat_by_chat_id(chat_id)
  243. shared_chat.title = chat.title
  244. shared_chat.chat = chat.chat
  245. shared_chat.updated_at = int(time.time())
  246. db.commit()
  247. db.refresh(shared_chat)
  248. return ChatModel.model_validate(shared_chat)
  249. except Exception:
  250. return None
  251. def delete_shared_chat_by_chat_id(self, chat_id: str) -> bool:
  252. try:
  253. with get_db() as db:
  254. db.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete()
  255. db.commit()
  256. return True
  257. except Exception:
  258. return False
  259. def update_chat_share_id_by_id(
  260. self, id: str, share_id: Optional[str]
  261. ) -> Optional[ChatModel]:
  262. try:
  263. with get_db() as db:
  264. chat = db.get(Chat, id)
  265. chat.share_id = share_id
  266. db.commit()
  267. db.refresh(chat)
  268. return ChatModel.model_validate(chat)
  269. except Exception:
  270. return None
  271. def toggle_chat_pinned_by_id(self, id: str) -> Optional[ChatModel]:
  272. try:
  273. with get_db() as db:
  274. chat = db.get(Chat, id)
  275. chat.pinned = not chat.pinned
  276. chat.updated_at = int(time.time())
  277. db.commit()
  278. db.refresh(chat)
  279. return ChatModel.model_validate(chat)
  280. except Exception:
  281. return None
  282. def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]:
  283. try:
  284. with get_db() as db:
  285. chat = db.get(Chat, id)
  286. chat.archived = not chat.archived
  287. chat.updated_at = int(time.time())
  288. db.commit()
  289. db.refresh(chat)
  290. return ChatModel.model_validate(chat)
  291. except Exception:
  292. return None
  293. def archive_all_chats_by_user_id(self, user_id: str) -> bool:
  294. try:
  295. with get_db() as db:
  296. db.query(Chat).filter_by(user_id=user_id).update({"archived": True})
  297. db.commit()
  298. return True
  299. except Exception:
  300. return False
  301. def get_archived_chat_list_by_user_id(
  302. self, user_id: str, skip: int = 0, limit: int = 50
  303. ) -> list[ChatModel]:
  304. with get_db() as db:
  305. all_chats = (
  306. db.query(Chat)
  307. .filter_by(user_id=user_id, archived=True)
  308. .order_by(Chat.updated_at.desc())
  309. # .limit(limit).offset(skip)
  310. .all()
  311. )
  312. return [ChatModel.model_validate(chat) for chat in all_chats]
  313. def get_chat_list_by_user_id(
  314. self,
  315. user_id: str,
  316. include_archived: bool = False,
  317. skip: int = 0,
  318. limit: int = 50,
  319. ) -> list[ChatModel]:
  320. with get_db() as db:
  321. query = db.query(Chat).filter_by(user_id=user_id).filter_by(folder_id=None)
  322. if not include_archived:
  323. query = query.filter_by(archived=False)
  324. query = query.order_by(Chat.updated_at.desc())
  325. if skip:
  326. query = query.offset(skip)
  327. if limit:
  328. query = query.limit(limit)
  329. all_chats = query.all()
  330. return [ChatModel.model_validate(chat) for chat in all_chats]
  331. def get_chat_title_id_list_by_user_id(
  332. self,
  333. user_id: str,
  334. include_archived: bool = False,
  335. skip: Optional[int] = None,
  336. limit: Optional[int] = None,
  337. ) -> list[ChatTitleIdResponse]:
  338. with get_db() as db:
  339. query = db.query(Chat).filter_by(user_id=user_id).filter_by(folder_id=None)
  340. query = query.filter(or_(Chat.pinned == False, Chat.pinned == None))
  341. if not include_archived:
  342. query = query.filter_by(archived=False)
  343. query = query.order_by(Chat.updated_at.desc()).with_entities(
  344. Chat.id, Chat.title, Chat.updated_at, Chat.created_at
  345. )
  346. if skip:
  347. query = query.offset(skip)
  348. if limit:
  349. query = query.limit(limit)
  350. all_chats = query.all()
  351. # result has to be destrctured from sqlalchemy `row` and mapped to a dict since the `ChatModel`is not the returned dataclass.
  352. return [
  353. ChatTitleIdResponse.model_validate(
  354. {
  355. "id": chat[0],
  356. "title": chat[1],
  357. "updated_at": chat[2],
  358. "created_at": chat[3],
  359. }
  360. )
  361. for chat in all_chats
  362. ]
  363. def get_chat_list_by_chat_ids(
  364. self, chat_ids: list[str], skip: int = 0, limit: int = 50
  365. ) -> list[ChatModel]:
  366. with get_db() as db:
  367. all_chats = (
  368. db.query(Chat)
  369. .filter(Chat.id.in_(chat_ids))
  370. .filter_by(archived=False)
  371. .order_by(Chat.updated_at.desc())
  372. .all()
  373. )
  374. return [ChatModel.model_validate(chat) for chat in all_chats]
  375. def get_chat_by_id(self, id: str) -> Optional[ChatModel]:
  376. try:
  377. with get_db() as db:
  378. chat = db.get(Chat, id)
  379. return ChatModel.model_validate(chat)
  380. except Exception:
  381. return None
  382. def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]:
  383. try:
  384. with get_db() as db:
  385. # it is possible that the shared link was deleted. hence,
  386. # we check if the chat is still shared by checkng if a chat with the share_id exists
  387. chat = db.query(Chat).filter_by(share_id=id).first()
  388. if chat:
  389. return self.get_chat_by_id(id)
  390. else:
  391. return None
  392. except Exception:
  393. return None
  394. def get_chat_by_id_and_user_id(self, id: str, user_id: str) -> Optional[ChatModel]:
  395. try:
  396. with get_db() as db:
  397. chat = db.query(Chat).filter_by(id=id, user_id=user_id).first()
  398. return ChatModel.model_validate(chat)
  399. except Exception:
  400. return None
  401. def get_chats(self, skip: int = 0, limit: int = 50) -> list[ChatModel]:
  402. with get_db() as db:
  403. all_chats = (
  404. db.query(Chat)
  405. # .limit(limit).offset(skip)
  406. .order_by(Chat.updated_at.desc())
  407. )
  408. return [ChatModel.model_validate(chat) for chat in all_chats]
  409. def get_chats_by_user_id(self, user_id: str) -> list[ChatModel]:
  410. with get_db() as db:
  411. all_chats = (
  412. db.query(Chat)
  413. .filter_by(user_id=user_id)
  414. .order_by(Chat.updated_at.desc())
  415. )
  416. return [ChatModel.model_validate(chat) for chat in all_chats]
  417. def get_pinned_chats_by_user_id(self, user_id: str) -> list[ChatModel]:
  418. with get_db() as db:
  419. all_chats = (
  420. db.query(Chat)
  421. .filter_by(user_id=user_id, pinned=True, archived=False)
  422. .order_by(Chat.updated_at.desc())
  423. )
  424. return [ChatModel.model_validate(chat) for chat in all_chats]
  425. def get_archived_chats_by_user_id(self, user_id: str) -> list[ChatModel]:
  426. with get_db() as db:
  427. all_chats = (
  428. db.query(Chat)
  429. .filter_by(user_id=user_id, archived=True)
  430. .order_by(Chat.updated_at.desc())
  431. )
  432. return [ChatModel.model_validate(chat) for chat in all_chats]
  433. def get_chats_by_user_id_and_search_text(
  434. self,
  435. user_id: str,
  436. search_text: str,
  437. include_archived: bool = False,
  438. skip: int = 0,
  439. limit: int = 60,
  440. ) -> list[ChatModel]:
  441. """
  442. Filters chats based on a search query using Python, allowing pagination using skip and limit.
  443. """
  444. search_text = search_text.lower().strip()
  445. if not search_text:
  446. return self.get_chat_list_by_user_id(user_id, include_archived, skip, limit)
  447. search_text_words = search_text.split(" ")
  448. # search_text might contain 'tag:tag_name' format so we need to extract the tag_name, split the search_text and remove the tags
  449. tag_ids = [
  450. word.replace("tag:", "").replace(" ", "_").lower()
  451. for word in search_text_words
  452. if word.startswith("tag:")
  453. ]
  454. search_text_words = [
  455. word for word in search_text_words if not word.startswith("tag:")
  456. ]
  457. search_text = " ".join(search_text_words)
  458. with get_db() as db:
  459. query = db.query(Chat).filter(Chat.user_id == user_id)
  460. if not include_archived:
  461. query = query.filter(Chat.archived == False)
  462. query = query.order_by(Chat.updated_at.desc())
  463. # Check if the database dialect is either 'sqlite' or 'postgresql'
  464. dialect_name = db.bind.dialect.name
  465. if dialect_name == "sqlite":
  466. # SQLite case: using JSON1 extension for JSON searching
  467. query = query.filter(
  468. (
  469. Chat.title.ilike(
  470. f"%{search_text}%"
  471. ) # Case-insensitive search in title
  472. | text(
  473. """
  474. EXISTS (
  475. SELECT 1
  476. FROM json_each(Chat.chat, '$.messages') AS message
  477. WHERE LOWER(message.value->>'content') LIKE '%' || :search_text || '%'
  478. )
  479. """
  480. )
  481. ).params(search_text=search_text)
  482. )
  483. # Check if there are any tags to filter, it should have all the tags
  484. if "none" in tag_ids:
  485. query = query.filter(
  486. text(
  487. """
  488. NOT EXISTS (
  489. SELECT 1
  490. FROM json_each(Chat.meta, '$.tags') AS tag
  491. )
  492. """
  493. )
  494. )
  495. elif tag_ids:
  496. query = query.filter(
  497. and_(
  498. *[
  499. text(
  500. f"""
  501. EXISTS (
  502. SELECT 1
  503. FROM json_each(Chat.meta, '$.tags') AS tag
  504. WHERE tag.value = :tag_id_{tag_idx}
  505. )
  506. """
  507. ).params(**{f"tag_id_{tag_idx}": tag_id})
  508. for tag_idx, tag_id in enumerate(tag_ids)
  509. ]
  510. )
  511. )
  512. elif dialect_name == "postgresql":
  513. # PostgreSQL relies on proper JSON query for search
  514. query = query.filter(
  515. (
  516. Chat.title.ilike(
  517. f"%{search_text}%"
  518. ) # Case-insensitive search in title
  519. | text(
  520. """
  521. EXISTS (
  522. SELECT 1
  523. FROM json_array_elements(Chat.chat->'messages') AS message
  524. WHERE LOWER(message->>'content') LIKE '%' || :search_text || '%'
  525. )
  526. """
  527. )
  528. ).params(search_text=search_text)
  529. )
  530. # Check if there are any tags to filter, it should have all the tags
  531. if "none" in tag_ids:
  532. query = query.filter(
  533. text(
  534. """
  535. NOT EXISTS (
  536. SELECT 1
  537. FROM json_array_elements_text(Chat.meta->'tags') AS tag
  538. )
  539. """
  540. )
  541. )
  542. elif tag_ids:
  543. query = query.filter(
  544. and_(
  545. *[
  546. text(
  547. f"""
  548. EXISTS (
  549. SELECT 1
  550. FROM json_array_elements_text(Chat.meta->'tags') AS tag
  551. WHERE tag = :tag_id_{tag_idx}
  552. )
  553. """
  554. ).params(**{f"tag_id_{tag_idx}": tag_id})
  555. for tag_idx, tag_id in enumerate(tag_ids)
  556. ]
  557. )
  558. )
  559. else:
  560. raise NotImplementedError(
  561. f"Unsupported dialect: {db.bind.dialect.name}"
  562. )
  563. # Perform pagination at the SQL level
  564. all_chats = query.offset(skip).limit(limit).all()
  565. print(len(all_chats))
  566. # Validate and return chats
  567. return [ChatModel.model_validate(chat) for chat in all_chats]
  568. def get_chats_by_folder_id_and_user_id(
  569. self, folder_id: str, user_id: str
  570. ) -> list[ChatModel]:
  571. with get_db() as db:
  572. query = db.query(Chat).filter_by(folder_id=folder_id, user_id=user_id)
  573. query = query.filter(or_(Chat.pinned == False, Chat.pinned == None))
  574. query = query.filter_by(archived=False)
  575. query = query.order_by(Chat.updated_at.desc())
  576. all_chats = query.all()
  577. return [ChatModel.model_validate(chat) for chat in all_chats]
  578. def get_chats_by_folder_ids_and_user_id(
  579. self, folder_ids: list[str], user_id: str
  580. ) -> list[ChatModel]:
  581. with get_db() as db:
  582. query = db.query(Chat).filter(
  583. Chat.folder_id.in_(folder_ids), Chat.user_id == user_id
  584. )
  585. query = query.filter(or_(Chat.pinned == False, Chat.pinned == None))
  586. query = query.filter_by(archived=False)
  587. query = query.order_by(Chat.updated_at.desc())
  588. all_chats = query.all()
  589. return [ChatModel.model_validate(chat) for chat in all_chats]
  590. def update_chat_folder_id_by_id_and_user_id(
  591. self, id: str, user_id: str, folder_id: str
  592. ) -> Optional[ChatModel]:
  593. try:
  594. with get_db() as db:
  595. chat = db.get(Chat, id)
  596. chat.folder_id = folder_id
  597. chat.updated_at = int(time.time())
  598. chat.pinned = False
  599. db.commit()
  600. db.refresh(chat)
  601. return ChatModel.model_validate(chat)
  602. except Exception:
  603. return None
  604. def get_chat_tags_by_id_and_user_id(self, id: str, user_id: str) -> list[TagModel]:
  605. with get_db() as db:
  606. chat = db.get(Chat, id)
  607. tags = chat.meta.get("tags", [])
  608. return [Tags.get_tag_by_name_and_user_id(tag, user_id) for tag in tags]
  609. def get_chat_list_by_user_id_and_tag_name(
  610. self, user_id: str, tag_name: str, skip: int = 0, limit: int = 50
  611. ) -> list[ChatModel]:
  612. with get_db() as db:
  613. query = db.query(Chat).filter_by(user_id=user_id)
  614. tag_id = tag_name.replace(" ", "_").lower()
  615. print(db.bind.dialect.name)
  616. if db.bind.dialect.name == "sqlite":
  617. # SQLite JSON1 querying for tags within the meta JSON field
  618. query = query.filter(
  619. text(
  620. f"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)"
  621. )
  622. ).params(tag_id=tag_id)
  623. elif db.bind.dialect.name == "postgresql":
  624. # PostgreSQL JSON query for tags within the meta JSON field (for `json` type)
  625. query = query.filter(
  626. text(
  627. "EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)"
  628. )
  629. ).params(tag_id=tag_id)
  630. else:
  631. raise NotImplementedError(
  632. f"Unsupported dialect: {db.bind.dialect.name}"
  633. )
  634. all_chats = query.all()
  635. print("all_chats", all_chats)
  636. return [ChatModel.model_validate(chat) for chat in all_chats]
  637. def add_chat_tag_by_id_and_user_id_and_tag_name(
  638. self, id: str, user_id: str, tag_name: str
  639. ) -> Optional[ChatModel]:
  640. tag = Tags.get_tag_by_name_and_user_id(tag_name, user_id)
  641. if tag is None:
  642. tag = Tags.insert_new_tag(tag_name, user_id)
  643. try:
  644. with get_db() as db:
  645. chat = db.get(Chat, id)
  646. tag_id = tag.id
  647. if tag_id not in chat.meta.get("tags", []):
  648. chat.meta = {
  649. **chat.meta,
  650. "tags": list(set(chat.meta.get("tags", []) + [tag_id])),
  651. }
  652. db.commit()
  653. db.refresh(chat)
  654. return ChatModel.model_validate(chat)
  655. except Exception:
  656. return None
  657. def count_chats_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> int:
  658. with get_db() as db: # Assuming `get_db()` returns a session object
  659. query = db.query(Chat).filter_by(user_id=user_id, archived=False)
  660. # Normalize the tag_name for consistency
  661. tag_id = tag_name.replace(" ", "_").lower()
  662. if db.bind.dialect.name == "sqlite":
  663. # SQLite JSON1 support for querying the tags inside the `meta` JSON field
  664. query = query.filter(
  665. text(
  666. f"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)"
  667. )
  668. ).params(tag_id=tag_id)
  669. elif db.bind.dialect.name == "postgresql":
  670. # PostgreSQL JSONB support for querying the tags inside the `meta` JSON field
  671. query = query.filter(
  672. text(
  673. "EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)"
  674. )
  675. ).params(tag_id=tag_id)
  676. else:
  677. raise NotImplementedError(
  678. f"Unsupported dialect: {db.bind.dialect.name}"
  679. )
  680. # Get the count of matching records
  681. count = query.count()
  682. # Debugging output for inspection
  683. print(f"Count of chats for tag '{tag_name}':", count)
  684. return count
  685. def delete_tag_by_id_and_user_id_and_tag_name(
  686. self, id: str, user_id: str, tag_name: str
  687. ) -> bool:
  688. try:
  689. with get_db() as db:
  690. chat = db.get(Chat, id)
  691. tags = chat.meta.get("tags", [])
  692. tag_id = tag_name.replace(" ", "_").lower()
  693. tags = [tag for tag in tags if tag != tag_id]
  694. chat.meta = {
  695. **chat.meta,
  696. "tags": list(set(tags)),
  697. }
  698. db.commit()
  699. return True
  700. except Exception:
  701. return False
  702. def delete_all_tags_by_id_and_user_id(self, id: str, user_id: str) -> bool:
  703. try:
  704. with get_db() as db:
  705. chat = db.get(Chat, id)
  706. chat.meta = {
  707. **chat.meta,
  708. "tags": [],
  709. }
  710. db.commit()
  711. return True
  712. except Exception:
  713. return False
  714. def delete_chat_by_id(self, id: str) -> bool:
  715. try:
  716. with get_db() as db:
  717. db.query(Chat).filter_by(id=id).delete()
  718. db.commit()
  719. return True and self.delete_shared_chat_by_chat_id(id)
  720. except Exception:
  721. return False
  722. def delete_chat_by_id_and_user_id(self, id: str, user_id: str) -> bool:
  723. try:
  724. with get_db() as db:
  725. db.query(Chat).filter_by(id=id, user_id=user_id).delete()
  726. db.commit()
  727. return True and self.delete_shared_chat_by_chat_id(id)
  728. except Exception:
  729. return False
  730. def delete_chats_by_user_id(self, user_id: str) -> bool:
  731. try:
  732. with get_db() as db:
  733. self.delete_shared_chats_by_user_id(user_id)
  734. db.query(Chat).filter_by(user_id=user_id).delete()
  735. db.commit()
  736. return True
  737. except Exception:
  738. return False
  739. def delete_chats_by_user_id_and_folder_id(
  740. self, user_id: str, folder_id: str
  741. ) -> bool:
  742. try:
  743. with get_db() as db:
  744. db.query(Chat).filter_by(user_id=user_id, folder_id=folder_id).delete()
  745. db.commit()
  746. return True
  747. except Exception:
  748. return False
  749. def delete_shared_chats_by_user_id(self, user_id: str) -> bool:
  750. try:
  751. with get_db() as db:
  752. chats_by_user = db.query(Chat).filter_by(user_id=user_id).all()
  753. shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]
  754. db.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete()
  755. db.commit()
  756. return True
  757. except Exception:
  758. return False
  759. Chats = ChatTable()