chats.py 32 KB

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