chats.py 37 KB

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