chats.py 36 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057
  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 update_chat_share_id_by_id(
  299. self, id: str, share_id: Optional[str]
  300. ) -> Optional[ChatModel]:
  301. try:
  302. with get_db() as db:
  303. chat = db.get(Chat, id)
  304. chat.share_id = share_id
  305. db.commit()
  306. db.refresh(chat)
  307. return ChatModel.model_validate(chat)
  308. except Exception:
  309. return None
  310. def toggle_chat_pinned_by_id(self, id: str) -> Optional[ChatModel]:
  311. try:
  312. with get_db() as db:
  313. chat = db.get(Chat, id)
  314. chat.pinned = not chat.pinned
  315. chat.updated_at = int(time.time())
  316. db.commit()
  317. db.refresh(chat)
  318. return ChatModel.model_validate(chat)
  319. except Exception:
  320. return None
  321. def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]:
  322. try:
  323. with get_db() as db:
  324. chat = db.get(Chat, id)
  325. chat.archived = not chat.archived
  326. chat.updated_at = int(time.time())
  327. db.commit()
  328. db.refresh(chat)
  329. return ChatModel.model_validate(chat)
  330. except Exception:
  331. return None
  332. def archive_all_chats_by_user_id(self, user_id: str) -> bool:
  333. try:
  334. with get_db() as db:
  335. db.query(Chat).filter_by(user_id=user_id).update({"archived": True})
  336. db.commit()
  337. return True
  338. except Exception:
  339. return False
  340. def get_archived_chat_list_by_user_id(
  341. self,
  342. user_id: str,
  343. filter: Optional[dict] = None,
  344. skip: int = 0,
  345. limit: int = 50,
  346. ) -> list[ChatModel]:
  347. with get_db() as db:
  348. query = db.query(Chat).filter_by(user_id=user_id, archived=True)
  349. if filter:
  350. query_key = filter.get("query")
  351. if query_key:
  352. query = query.filter(Chat.title.ilike(f"%{query_key}%"))
  353. order_by = filter.get("order_by")
  354. direction = filter.get("direction")
  355. if order_by and direction and getattr(Chat, order_by):
  356. if direction.lower() == "asc":
  357. query = query.order_by(getattr(Chat, order_by).asc())
  358. elif direction.lower() == "desc":
  359. query = query.order_by(getattr(Chat, order_by).desc())
  360. else:
  361. raise ValueError("Invalid direction for ordering")
  362. else:
  363. query = query.order_by(Chat.updated_at.desc())
  364. if skip:
  365. query = query.offset(skip)
  366. if limit:
  367. query = query.limit(limit)
  368. all_chats = query.all()
  369. return [ChatModel.model_validate(chat) for chat in all_chats]
  370. def get_chat_list_by_user_id(
  371. self,
  372. user_id: str,
  373. include_archived: bool = False,
  374. filter: Optional[dict] = None,
  375. skip: int = 0,
  376. limit: int = 50,
  377. ) -> list[ChatModel]:
  378. with get_db() as db:
  379. query = db.query(Chat).filter_by(user_id=user_id)
  380. if not include_archived:
  381. query = query.filter_by(archived=False)
  382. if filter:
  383. query_key = filter.get("query")
  384. if query_key:
  385. query = query.filter(Chat.title.ilike(f"%{query_key}%"))
  386. order_by = filter.get("order_by")
  387. direction = filter.get("direction")
  388. if order_by and direction and getattr(Chat, order_by):
  389. if direction.lower() == "asc":
  390. query = query.order_by(getattr(Chat, order_by).asc())
  391. elif direction.lower() == "desc":
  392. query = query.order_by(getattr(Chat, order_by).desc())
  393. else:
  394. raise ValueError("Invalid direction for ordering")
  395. else:
  396. query = query.order_by(Chat.updated_at.desc())
  397. if skip:
  398. query = query.offset(skip)
  399. if limit:
  400. query = query.limit(limit)
  401. all_chats = query.all()
  402. return [ChatModel.model_validate(chat) for chat in all_chats]
  403. def get_chat_title_id_list_by_user_id(
  404. self,
  405. user_id: str,
  406. include_archived: bool = False,
  407. include_folders: bool = False,
  408. skip: Optional[int] = None,
  409. limit: Optional[int] = None,
  410. ) -> list[ChatTitleIdResponse]:
  411. with get_db() as db:
  412. query = db.query(Chat).filter_by(user_id=user_id)
  413. if not include_folders:
  414. query = query.filter_by(folder_id=None)
  415. query = query.filter(or_(Chat.pinned == False, Chat.pinned == None))
  416. if not include_archived:
  417. query = query.filter_by(archived=False)
  418. query = query.order_by(Chat.updated_at.desc()).with_entities(
  419. Chat.id, Chat.title, Chat.updated_at, Chat.created_at
  420. )
  421. if skip:
  422. query = query.offset(skip)
  423. if limit:
  424. query = query.limit(limit)
  425. all_chats = query.all()
  426. # result has to be destructured from sqlalchemy `row` and mapped to a dict since the `ChatModel`is not the returned dataclass.
  427. return [
  428. ChatTitleIdResponse.model_validate(
  429. {
  430. "id": chat[0],
  431. "title": chat[1],
  432. "updated_at": chat[2],
  433. "created_at": chat[3],
  434. }
  435. )
  436. for chat in all_chats
  437. ]
  438. def get_chat_list_by_chat_ids(
  439. self, chat_ids: list[str], skip: int = 0, limit: int = 50
  440. ) -> list[ChatModel]:
  441. with get_db() as db:
  442. all_chats = (
  443. db.query(Chat)
  444. .filter(Chat.id.in_(chat_ids))
  445. .filter_by(archived=False)
  446. .order_by(Chat.updated_at.desc())
  447. .all()
  448. )
  449. return [ChatModel.model_validate(chat) for chat in all_chats]
  450. def get_chat_by_id(self, id: str) -> Optional[ChatModel]:
  451. try:
  452. with get_db() as db:
  453. chat = db.get(Chat, id)
  454. return ChatModel.model_validate(chat)
  455. except Exception:
  456. return None
  457. def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]:
  458. try:
  459. with get_db() as db:
  460. # it is possible that the shared link was deleted. hence,
  461. # we check if the chat is still shared by checking if a chat with the share_id exists
  462. chat = db.query(Chat).filter_by(share_id=id).first()
  463. if chat:
  464. return self.get_chat_by_id(id)
  465. else:
  466. return None
  467. except Exception:
  468. return None
  469. def get_chat_by_id_and_user_id(self, id: str, user_id: str) -> Optional[ChatModel]:
  470. try:
  471. with get_db() as db:
  472. chat = db.query(Chat).filter_by(id=id, user_id=user_id).first()
  473. return ChatModel.model_validate(chat)
  474. except Exception:
  475. return None
  476. def get_chats(self, skip: int = 0, limit: int = 50) -> list[ChatModel]:
  477. with get_db() as db:
  478. all_chats = (
  479. db.query(Chat)
  480. # .limit(limit).offset(skip)
  481. .order_by(Chat.updated_at.desc())
  482. )
  483. return [ChatModel.model_validate(chat) for chat in all_chats]
  484. def get_chats_by_user_id(self, user_id: str) -> list[ChatModel]:
  485. with get_db() as db:
  486. all_chats = (
  487. db.query(Chat)
  488. .filter_by(user_id=user_id)
  489. .order_by(Chat.updated_at.desc())
  490. )
  491. return [ChatModel.model_validate(chat) for chat in all_chats]
  492. def get_pinned_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, pinned=True, archived=False)
  497. .order_by(Chat.updated_at.desc())
  498. )
  499. return [ChatModel.model_validate(chat) for chat in all_chats]
  500. def get_archived_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, archived=True)
  505. .order_by(Chat.updated_at.desc())
  506. )
  507. return [ChatModel.model_validate(chat) for chat in all_chats]
  508. def get_chats_by_user_id_and_search_text(
  509. self,
  510. user_id: str,
  511. search_text: str,
  512. include_archived: bool = False,
  513. skip: int = 0,
  514. limit: int = 60,
  515. ) -> list[ChatModel]:
  516. """
  517. Filters chats based on a search query using Python, allowing pagination using skip and limit.
  518. """
  519. search_text = search_text.replace("\u0000", "").lower().strip()
  520. if not search_text:
  521. return self.get_chat_list_by_user_id(
  522. user_id, include_archived, filter={}, skip=skip, limit=limit
  523. )
  524. search_text_words = search_text.split(" ")
  525. # search_text might contain 'tag:tag_name' format so we need to extract the tag_name, split the search_text and remove the tags
  526. tag_ids = [
  527. word.replace("tag:", "").replace(" ", "_").lower()
  528. for word in search_text_words
  529. if word.startswith("tag:")
  530. ]
  531. # Extract folder names - handle spaces and case insensitivity
  532. folders = Folders.search_folders_by_names(
  533. user_id,
  534. [
  535. word.replace("folder:", "")
  536. for word in search_text_words
  537. if word.startswith("folder:")
  538. ],
  539. )
  540. folder_ids = [folder.id for folder in folders]
  541. is_pinned = None
  542. if "pinned:true" in search_text_words:
  543. is_pinned = True
  544. elif "pinned:false" in search_text_words:
  545. is_pinned = False
  546. is_archived = None
  547. if "archived:true" in search_text_words:
  548. is_archived = True
  549. elif "archived:false" in search_text_words:
  550. is_archived = False
  551. is_shared = None
  552. if "shared:true" in search_text_words:
  553. is_shared = True
  554. elif "shared:false" in search_text_words:
  555. is_shared = False
  556. search_text_words = [
  557. word
  558. for word in search_text_words
  559. if (
  560. not word.startswith("tag:")
  561. and not word.startswith("folder:")
  562. and not word.startswith("pinned:")
  563. and not word.startswith("archived:")
  564. and not word.startswith("shared:")
  565. )
  566. ]
  567. search_text = " ".join(search_text_words)
  568. with get_db() as db:
  569. query = db.query(Chat).filter(Chat.user_id == user_id)
  570. if is_archived is not None:
  571. query = query.filter(Chat.archived == is_archived)
  572. elif not include_archived:
  573. query = query.filter(Chat.archived == False)
  574. if is_pinned is not None:
  575. query = query.filter(Chat.pinned == is_pinned)
  576. if is_shared is not None:
  577. if is_shared:
  578. query = query.filter(Chat.share_id.isnot(None))
  579. else:
  580. query = query.filter(Chat.share_id.is_(None))
  581. if folder_ids:
  582. query = query.filter(Chat.folder_id.in_(folder_ids))
  583. query = query.order_by(Chat.updated_at.desc())
  584. # Check if the database dialect is either 'sqlite' or 'postgresql'
  585. dialect_name = db.bind.dialect.name
  586. if dialect_name == "sqlite":
  587. # SQLite case: using JSON1 extension for JSON searching
  588. sqlite_content_sql = (
  589. "EXISTS ("
  590. " SELECT 1 "
  591. " FROM json_each(Chat.chat, '$.messages') AS message "
  592. " WHERE LOWER(message.value->>'content') LIKE '%' || :content_key || '%'"
  593. ")"
  594. )
  595. sqlite_content_clause = text(sqlite_content_sql)
  596. query = query.filter(
  597. or_(
  598. Chat.title.ilike(bindparam("title_key")), sqlite_content_clause
  599. ).params(title_key=f"%{search_text}%", content_key=search_text)
  600. )
  601. # Check if there are any tags to filter, it should have all the tags
  602. if "none" in tag_ids:
  603. query = query.filter(
  604. text(
  605. """
  606. NOT EXISTS (
  607. SELECT 1
  608. FROM json_each(Chat.meta, '$.tags') AS tag
  609. )
  610. """
  611. )
  612. )
  613. elif tag_ids:
  614. query = query.filter(
  615. and_(
  616. *[
  617. text(
  618. f"""
  619. EXISTS (
  620. SELECT 1
  621. FROM json_each(Chat.meta, '$.tags') AS tag
  622. WHERE tag.value = :tag_id_{tag_idx}
  623. )
  624. """
  625. ).params(**{f"tag_id_{tag_idx}": tag_id})
  626. for tag_idx, tag_id in enumerate(tag_ids)
  627. ]
  628. )
  629. )
  630. elif dialect_name == "postgresql":
  631. # PostgreSQL relies on proper JSON query for search
  632. postgres_content_sql = (
  633. "EXISTS ("
  634. " SELECT 1 "
  635. " FROM json_array_elements(Chat.chat->'messages') AS message "
  636. " WHERE LOWER(message->>'content') LIKE '%' || :content_key || '%'"
  637. ")"
  638. )
  639. postgres_content_clause = text(postgres_content_sql)
  640. query = query.filter(
  641. or_(
  642. Chat.title.ilike(bindparam("title_key")),
  643. postgres_content_clause,
  644. ).params(title_key=f"%{search_text}%", content_key=search_text)
  645. )
  646. # Check if there are any tags to filter, it should have all the tags
  647. if "none" in tag_ids:
  648. query = query.filter(
  649. text(
  650. """
  651. NOT EXISTS (
  652. SELECT 1
  653. FROM json_array_elements_text(Chat.meta->'tags') AS tag
  654. )
  655. """
  656. )
  657. )
  658. elif tag_ids:
  659. query = query.filter(
  660. and_(
  661. *[
  662. text(
  663. f"""
  664. EXISTS (
  665. SELECT 1
  666. FROM json_array_elements_text(Chat.meta->'tags') AS tag
  667. WHERE tag = :tag_id_{tag_idx}
  668. )
  669. """
  670. ).params(**{f"tag_id_{tag_idx}": tag_id})
  671. for tag_idx, tag_id in enumerate(tag_ids)
  672. ]
  673. )
  674. )
  675. else:
  676. raise NotImplementedError(
  677. f"Unsupported dialect: {db.bind.dialect.name}"
  678. )
  679. # Perform pagination at the SQL level
  680. all_chats = query.offset(skip).limit(limit).all()
  681. log.info(f"The number of chats: {len(all_chats)}")
  682. # Validate and return chats
  683. return [ChatModel.model_validate(chat) for chat in all_chats]
  684. def get_chats_by_folder_id_and_user_id(
  685. self, folder_id: str, user_id: str, skip: int = 0, limit: int = 60
  686. ) -> list[ChatModel]:
  687. with get_db() as db:
  688. query = db.query(Chat).filter_by(folder_id=folder_id, user_id=user_id)
  689. query = query.filter(or_(Chat.pinned == False, Chat.pinned == None))
  690. query = query.filter_by(archived=False)
  691. query = query.order_by(Chat.updated_at.desc())
  692. if skip:
  693. query = query.offset(skip)
  694. if limit:
  695. query = query.limit(limit)
  696. all_chats = query.all()
  697. return [ChatModel.model_validate(chat) for chat in all_chats]
  698. def get_chats_by_folder_ids_and_user_id(
  699. self, folder_ids: list[str], user_id: str
  700. ) -> list[ChatModel]:
  701. with get_db() as db:
  702. query = db.query(Chat).filter(
  703. Chat.folder_id.in_(folder_ids), Chat.user_id == user_id
  704. )
  705. query = query.filter(or_(Chat.pinned == False, Chat.pinned == None))
  706. query = query.filter_by(archived=False)
  707. query = query.order_by(Chat.updated_at.desc())
  708. all_chats = query.all()
  709. return [ChatModel.model_validate(chat) for chat in all_chats]
  710. def update_chat_folder_id_by_id_and_user_id(
  711. self, id: str, user_id: str, folder_id: str
  712. ) -> Optional[ChatModel]:
  713. try:
  714. with get_db() as db:
  715. chat = db.get(Chat, id)
  716. chat.folder_id = folder_id
  717. chat.updated_at = int(time.time())
  718. chat.pinned = False
  719. db.commit()
  720. db.refresh(chat)
  721. return ChatModel.model_validate(chat)
  722. except Exception:
  723. return None
  724. def get_chat_tags_by_id_and_user_id(self, id: str, user_id: str) -> list[TagModel]:
  725. with get_db() as db:
  726. chat = db.get(Chat, id)
  727. tags = chat.meta.get("tags", [])
  728. return [Tags.get_tag_by_name_and_user_id(tag, user_id) for tag in tags]
  729. def get_chat_list_by_user_id_and_tag_name(
  730. self, user_id: str, tag_name: str, skip: int = 0, limit: int = 50
  731. ) -> list[ChatModel]:
  732. with get_db() as db:
  733. query = db.query(Chat).filter_by(user_id=user_id)
  734. tag_id = tag_name.replace(" ", "_").lower()
  735. log.info(f"DB dialect name: {db.bind.dialect.name}")
  736. if db.bind.dialect.name == "sqlite":
  737. # SQLite JSON1 querying for tags within the meta JSON field
  738. query = query.filter(
  739. text(
  740. f"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)"
  741. )
  742. ).params(tag_id=tag_id)
  743. elif db.bind.dialect.name == "postgresql":
  744. # PostgreSQL JSON query for tags within the meta JSON field (for `json` type)
  745. query = query.filter(
  746. text(
  747. "EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)"
  748. )
  749. ).params(tag_id=tag_id)
  750. else:
  751. raise NotImplementedError(
  752. f"Unsupported dialect: {db.bind.dialect.name}"
  753. )
  754. all_chats = query.all()
  755. log.debug(f"all_chats: {all_chats}")
  756. return [ChatModel.model_validate(chat) for chat in all_chats]
  757. def add_chat_tag_by_id_and_user_id_and_tag_name(
  758. self, id: str, user_id: str, tag_name: str
  759. ) -> Optional[ChatModel]:
  760. tag = Tags.get_tag_by_name_and_user_id(tag_name, user_id)
  761. if tag is None:
  762. tag = Tags.insert_new_tag(tag_name, user_id)
  763. try:
  764. with get_db() as db:
  765. chat = db.get(Chat, id)
  766. tag_id = tag.id
  767. if tag_id not in chat.meta.get("tags", []):
  768. chat.meta = {
  769. **chat.meta,
  770. "tags": list(set(chat.meta.get("tags", []) + [tag_id])),
  771. }
  772. db.commit()
  773. db.refresh(chat)
  774. return ChatModel.model_validate(chat)
  775. except Exception:
  776. return None
  777. def count_chats_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> int:
  778. with get_db() as db: # Assuming `get_db()` returns a session object
  779. query = db.query(Chat).filter_by(user_id=user_id, archived=False)
  780. # Normalize the tag_name for consistency
  781. tag_id = tag_name.replace(" ", "_").lower()
  782. if db.bind.dialect.name == "sqlite":
  783. # SQLite JSON1 support for querying the tags inside the `meta` JSON field
  784. query = query.filter(
  785. text(
  786. f"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)"
  787. )
  788. ).params(tag_id=tag_id)
  789. elif db.bind.dialect.name == "postgresql":
  790. # PostgreSQL JSONB support for querying the tags inside the `meta` JSON field
  791. query = query.filter(
  792. text(
  793. "EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)"
  794. )
  795. ).params(tag_id=tag_id)
  796. else:
  797. raise NotImplementedError(
  798. f"Unsupported dialect: {db.bind.dialect.name}"
  799. )
  800. # Get the count of matching records
  801. count = query.count()
  802. # Debugging output for inspection
  803. log.info(f"Count of chats for tag '{tag_name}': {count}")
  804. return count
  805. def count_chats_by_folder_id_and_user_id(self, folder_id: str, user_id: str) -> int:
  806. with get_db() as db:
  807. query = db.query(Chat).filter_by(user_id=user_id)
  808. query = query.filter_by(folder_id=folder_id)
  809. count = query.count()
  810. log.info(f"Count of chats for folder '{folder_id}': {count}")
  811. return count
  812. def delete_tag_by_id_and_user_id_and_tag_name(
  813. self, id: str, user_id: str, tag_name: str
  814. ) -> bool:
  815. try:
  816. with get_db() as db:
  817. chat = db.get(Chat, id)
  818. tags = chat.meta.get("tags", [])
  819. tag_id = tag_name.replace(" ", "_").lower()
  820. tags = [tag for tag in tags if tag != tag_id]
  821. chat.meta = {
  822. **chat.meta,
  823. "tags": list(set(tags)),
  824. }
  825. db.commit()
  826. return True
  827. except Exception:
  828. return False
  829. def delete_all_tags_by_id_and_user_id(self, id: str, user_id: str) -> bool:
  830. try:
  831. with get_db() as db:
  832. chat = db.get(Chat, id)
  833. chat.meta = {
  834. **chat.meta,
  835. "tags": [],
  836. }
  837. db.commit()
  838. return True
  839. except Exception:
  840. return False
  841. def delete_chat_by_id(self, id: str) -> bool:
  842. try:
  843. with get_db() as db:
  844. db.query(Chat).filter_by(id=id).delete()
  845. db.commit()
  846. return True and self.delete_shared_chat_by_chat_id(id)
  847. except Exception:
  848. return False
  849. def delete_chat_by_id_and_user_id(self, id: str, user_id: str) -> bool:
  850. try:
  851. with get_db() as db:
  852. db.query(Chat).filter_by(id=id, user_id=user_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_chats_by_user_id(self, user_id: str) -> bool:
  858. try:
  859. with get_db() as db:
  860. self.delete_shared_chats_by_user_id(user_id)
  861. db.query(Chat).filter_by(user_id=user_id).delete()
  862. db.commit()
  863. return True
  864. except Exception:
  865. return False
  866. def delete_chats_by_user_id_and_folder_id(
  867. self, user_id: str, folder_id: str
  868. ) -> bool:
  869. try:
  870. with get_db() as db:
  871. db.query(Chat).filter_by(user_id=user_id, folder_id=folder_id).delete()
  872. db.commit()
  873. return True
  874. except Exception:
  875. return False
  876. def delete_shared_chats_by_user_id(self, user_id: str) -> bool:
  877. try:
  878. with get_db() as db:
  879. chats_by_user = db.query(Chat).filter_by(user_id=user_id).all()
  880. shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]
  881. db.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete()
  882. db.commit()
  883. return True
  884. except Exception:
  885. return False
  886. Chats = ChatTable()