chats.py 36 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052
  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
  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. all_chats = query.all()
  693. return [ChatModel.model_validate(chat) for chat in all_chats]
  694. def get_chats_by_folder_ids_and_user_id(
  695. self, folder_ids: list[str], user_id: str
  696. ) -> list[ChatModel]:
  697. with get_db() as db:
  698. query = db.query(Chat).filter(
  699. Chat.folder_id.in_(folder_ids), Chat.user_id == user_id
  700. )
  701. query = query.filter(or_(Chat.pinned == False, Chat.pinned == None))
  702. query = query.filter_by(archived=False)
  703. query = query.order_by(Chat.updated_at.desc())
  704. all_chats = query.all()
  705. return [ChatModel.model_validate(chat) for chat in all_chats]
  706. def update_chat_folder_id_by_id_and_user_id(
  707. self, id: str, user_id: str, folder_id: str
  708. ) -> Optional[ChatModel]:
  709. try:
  710. with get_db() as db:
  711. chat = db.get(Chat, id)
  712. chat.folder_id = folder_id
  713. chat.updated_at = int(time.time())
  714. chat.pinned = False
  715. db.commit()
  716. db.refresh(chat)
  717. return ChatModel.model_validate(chat)
  718. except Exception:
  719. return None
  720. def get_chat_tags_by_id_and_user_id(self, id: str, user_id: str) -> list[TagModel]:
  721. with get_db() as db:
  722. chat = db.get(Chat, id)
  723. tags = chat.meta.get("tags", [])
  724. return [Tags.get_tag_by_name_and_user_id(tag, user_id) for tag in tags]
  725. def get_chat_list_by_user_id_and_tag_name(
  726. self, user_id: str, tag_name: str, skip: int = 0, limit: int = 50
  727. ) -> list[ChatModel]:
  728. with get_db() as db:
  729. query = db.query(Chat).filter_by(user_id=user_id)
  730. tag_id = tag_name.replace(" ", "_").lower()
  731. log.info(f"DB dialect name: {db.bind.dialect.name}")
  732. if db.bind.dialect.name == "sqlite":
  733. # SQLite JSON1 querying for tags within the meta JSON field
  734. query = query.filter(
  735. text(
  736. f"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)"
  737. )
  738. ).params(tag_id=tag_id)
  739. elif db.bind.dialect.name == "postgresql":
  740. # PostgreSQL JSON query for tags within the meta JSON field (for `json` type)
  741. query = query.filter(
  742. text(
  743. "EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)"
  744. )
  745. ).params(tag_id=tag_id)
  746. else:
  747. raise NotImplementedError(
  748. f"Unsupported dialect: {db.bind.dialect.name}"
  749. )
  750. all_chats = query.all()
  751. log.debug(f"all_chats: {all_chats}")
  752. return [ChatModel.model_validate(chat) for chat in all_chats]
  753. def add_chat_tag_by_id_and_user_id_and_tag_name(
  754. self, id: str, user_id: str, tag_name: str
  755. ) -> Optional[ChatModel]:
  756. tag = Tags.get_tag_by_name_and_user_id(tag_name, user_id)
  757. if tag is None:
  758. tag = Tags.insert_new_tag(tag_name, user_id)
  759. try:
  760. with get_db() as db:
  761. chat = db.get(Chat, id)
  762. tag_id = tag.id
  763. if tag_id not in chat.meta.get("tags", []):
  764. chat.meta = {
  765. **chat.meta,
  766. "tags": list(set(chat.meta.get("tags", []) + [tag_id])),
  767. }
  768. db.commit()
  769. db.refresh(chat)
  770. return ChatModel.model_validate(chat)
  771. except Exception:
  772. return None
  773. def count_chats_by_tag_name_and_user_id(self, tag_name: str, user_id: str) -> int:
  774. with get_db() as db: # Assuming `get_db()` returns a session object
  775. query = db.query(Chat).filter_by(user_id=user_id, archived=False)
  776. # Normalize the tag_name for consistency
  777. tag_id = tag_name.replace(" ", "_").lower()
  778. if db.bind.dialect.name == "sqlite":
  779. # SQLite JSON1 support for querying the tags inside the `meta` JSON field
  780. query = query.filter(
  781. text(
  782. f"EXISTS (SELECT 1 FROM json_each(Chat.meta, '$.tags') WHERE json_each.value = :tag_id)"
  783. )
  784. ).params(tag_id=tag_id)
  785. elif db.bind.dialect.name == "postgresql":
  786. # PostgreSQL JSONB support for querying the tags inside the `meta` JSON field
  787. query = query.filter(
  788. text(
  789. "EXISTS (SELECT 1 FROM json_array_elements_text(Chat.meta->'tags') elem WHERE elem = :tag_id)"
  790. )
  791. ).params(tag_id=tag_id)
  792. else:
  793. raise NotImplementedError(
  794. f"Unsupported dialect: {db.bind.dialect.name}"
  795. )
  796. # Get the count of matching records
  797. count = query.count()
  798. # Debugging output for inspection
  799. log.info(f"Count of chats for tag '{tag_name}': {count}")
  800. return count
  801. def count_chats_by_folder_id_and_user_id(self, folder_id: str, user_id: str) -> int:
  802. with get_db() as db:
  803. query = db.query(Chat).filter_by(user_id=user_id)
  804. query = query.filter_by(folder_id=folder_id)
  805. count = query.count()
  806. log.info(f"Count of chats for folder '{folder_id}': {count}")
  807. return count
  808. def delete_tag_by_id_and_user_id_and_tag_name(
  809. self, id: str, user_id: str, tag_name: str
  810. ) -> bool:
  811. try:
  812. with get_db() as db:
  813. chat = db.get(Chat, id)
  814. tags = chat.meta.get("tags", [])
  815. tag_id = tag_name.replace(" ", "_").lower()
  816. tags = [tag for tag in tags if tag != tag_id]
  817. chat.meta = {
  818. **chat.meta,
  819. "tags": list(set(tags)),
  820. }
  821. db.commit()
  822. return True
  823. except Exception:
  824. return False
  825. def delete_all_tags_by_id_and_user_id(self, id: str, user_id: str) -> bool:
  826. try:
  827. with get_db() as db:
  828. chat = db.get(Chat, id)
  829. chat.meta = {
  830. **chat.meta,
  831. "tags": [],
  832. }
  833. db.commit()
  834. return True
  835. except Exception:
  836. return False
  837. def delete_chat_by_id(self, id: str) -> bool:
  838. try:
  839. with get_db() as db:
  840. db.query(Chat).filter_by(id=id).delete()
  841. db.commit()
  842. return True and self.delete_shared_chat_by_chat_id(id)
  843. except Exception:
  844. return False
  845. def delete_chat_by_id_and_user_id(self, id: str, user_id: str) -> bool:
  846. try:
  847. with get_db() as db:
  848. db.query(Chat).filter_by(id=id, user_id=user_id).delete()
  849. db.commit()
  850. return True and self.delete_shared_chat_by_chat_id(id)
  851. except Exception:
  852. return False
  853. def delete_chats_by_user_id(self, user_id: str) -> bool:
  854. try:
  855. with get_db() as db:
  856. self.delete_shared_chats_by_user_id(user_id)
  857. db.query(Chat).filter_by(user_id=user_id).delete()
  858. db.commit()
  859. return True
  860. except Exception:
  861. return False
  862. def delete_chats_by_user_id_and_folder_id(
  863. self, user_id: str, folder_id: str
  864. ) -> bool:
  865. try:
  866. with get_db() as db:
  867. db.query(Chat).filter_by(user_id=user_id, folder_id=folder_id).delete()
  868. db.commit()
  869. return True
  870. except Exception:
  871. return False
  872. def delete_shared_chats_by_user_id(self, user_id: str) -> bool:
  873. try:
  874. with get_db() as db:
  875. chats_by_user = db.query(Chat).filter_by(user_id=user_id).all()
  876. shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]
  877. db.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete()
  878. db.commit()
  879. return True
  880. except Exception:
  881. return False
  882. Chats = ChatTable()