__init__.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382
  1. from app.schemas import UserCreate, PersonaCreate, PersonaUpdate, ForumCreate, MessageCreate
  2. from app.core.hashing import Hasher
  3. from app.db.client import fetch_one, fetch_all, RowObject, db_transaction, db_execute_commit
  4. from app.core.cache import cache_service
  5. import json
  6. import logging
  7. from typing import List, Optional, Any
  8. from datetime import datetime
  9. logger = logging.getLogger(__name__)
  10. def _normalize_persona(persona):
  11. if persona and isinstance(getattr(persona, "theories", None), str):
  12. try:
  13. theories = json.loads(persona.theories)
  14. if isinstance(theories, list):
  15. persona.theories = theories
  16. except json.JSONDecodeError:
  17. pass
  18. return persona
  19. # --- Cache Keys ---
  20. def user_cache_key(username: str): return f"user:{username}"
  21. def persona_cache_key(pid: int): return f"persona:{pid}"
  22. def forum_cache_key(fid: int): return f"forum:{fid}"
  23. def forum_participants_cache_key(fid: int): return f"forum:{fid}:participants"
  24. # --- User ---
  25. def get_user_by_username(db, username: str):
  26. # Cache Aside: Read
  27. cache_key = user_cache_key(username)
  28. cached = cache_service.get_cache(cache_key)
  29. if cached:
  30. return RowObject(cached) # Convert dict back to RowObject-like
  31. rs = db.execute("SELECT * FROM users WHERE username = ?", [username])
  32. user = fetch_one(rs)
  33. if user:
  34. cache_service.set_cache(cache_key, user.__dict__, expire=3600)
  35. return user
  36. def create_user(db: Any, user: UserCreate):
  37. password_bytes = user.password.encode('utf-8')
  38. if len(password_bytes) > 71:
  39. password_bytes = password_bytes[:71]
  40. safe_password = password_bytes.decode('utf-8', 'ignore')
  41. try:
  42. # Use transaction to ensure commit
  43. pwd_hash = Hasher.get_password_hash(safe_password)
  44. created_at = datetime.now()
  45. rs = db_execute_commit(
  46. db,
  47. "INSERT INTO users (username, email, password_hash, role, created_at) VALUES (?, ?, ?, ?, ?) RETURNING *",
  48. [user.username, user.email, pwd_hash, user.role, created_at]
  49. )
  50. new_user = fetch_one(rs)
  51. if new_user:
  52. cache_service.set_cache(user_cache_key(new_user.username), new_user.__dict__, expire=3600)
  53. return new_user
  54. except Exception as e:
  55. logger.error(f"Error creating user: {e}")
  56. raise
  57. # --- Persona ---
  58. def create_persona(db, persona: PersonaCreate, owner_id: int):
  59. try:
  60. theories_json = json.dumps(persona.theories)
  61. created_at = datetime.now()
  62. rs = db_execute_commit(
  63. db,
  64. """
  65. INSERT INTO personas (owner_id, name, title, bio, theories, stance, system_prompt, is_public, created_at)
  66. VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
  67. RETURNING *
  68. """,
  69. [
  70. owner_id,
  71. persona.name,
  72. persona.title,
  73. persona.bio,
  74. theories_json,
  75. persona.stance,
  76. persona.system_prompt,
  77. persona.is_public,
  78. created_at
  79. ]
  80. )
  81. new_persona = fetch_one(rs)
  82. # Cache Aside: Don't set cache on create. Let the first read populate it.
  83. # This ensures strict adherence to "DB is source of truth" and lazy loading.
  84. return _normalize_persona(new_persona)
  85. except Exception as e:
  86. logger.error(f"Error creating persona: {e}")
  87. raise
  88. def get_persona(db, persona_id: int):
  89. cache_key = persona_cache_key(persona_id)
  90. cached = cache_service.get_cache(cache_key)
  91. if cached:
  92. return _normalize_persona(RowObject(cached))
  93. rs = db.execute("SELECT * FROM personas WHERE id = ?", [persona_id])
  94. persona = fetch_one(rs)
  95. persona = _normalize_persona(persona)
  96. if persona:
  97. cache_service.set_cache(cache_key, persona.__dict__)
  98. return persona
  99. def update_persona(db, persona_id: int, updates: PersonaUpdate):
  100. try:
  101. update_data = updates.model_dump(exclude_unset=True)
  102. if not update_data:
  103. return get_persona(db, persona_id)
  104. set_clauses = []
  105. values = []
  106. for key, value in update_data.items():
  107. set_clauses.append(f"{key} = ?")
  108. if key == "theories":
  109. values.append(json.dumps(value))
  110. else:
  111. values.append(value)
  112. values.append(persona_id)
  113. query = f"UPDATE personas SET {', '.join(set_clauses)} WHERE id = ? RETURNING *"
  114. rs = db_execute_commit(db, query, values)
  115. updated = fetch_one(rs)
  116. # Sync Strategy: Delete Redis Key on Update
  117. if updated:
  118. cache_service.delete_cache(persona_cache_key(persona_id))
  119. return _normalize_persona(updated)
  120. except Exception as e:
  121. logger.error(f"Error updating persona: {e}")
  122. raise
  123. def delete_persona(db, persona_id: int):
  124. try:
  125. # Check if exists first to ensure idempotency and clear error
  126. rs_check = db.execute("SELECT id FROM personas WHERE id = ?", [persona_id])
  127. if not fetch_one(rs_check):
  128. return True # Already deleted or not exists
  129. with db_transaction(db) as tx:
  130. # Manually set persona_id to NULL in messages to avoid FK violation
  131. tx.execute("UPDATE messages SET persona_id = NULL WHERE persona_id = ?", [persona_id])
  132. # Cascading deletes should be handled by DB foreign keys,
  133. # but let's be explicit if needed or just execute
  134. rs = tx.execute("DELETE FROM personas WHERE id = ?", [persona_id])
  135. # FORCE COMMIT
  136. if hasattr(tx, 'commit'):
  137. tx.commit()
  138. elif hasattr(db, 'commit'):
  139. db.commit()
  140. # Sync Strategy: Delete Redis Key on Delete
  141. cache_service.delete_cache(persona_cache_key(persona_id))
  142. return True
  143. except Exception as e:
  144. logger.error(f"Error deleting persona {persona_id}: {e}")
  145. raise
  146. # --- Forum ---
  147. def create_forum(db, forum: ForumCreate, creator_id: int):
  148. try:
  149. with db_transaction(db) as tx:
  150. rs = tx.execute(
  151. """
  152. INSERT INTO forums (topic, creator_id, moderator_id, status, duration_minutes, start_time, summary_history, ablation_flags)
  153. VALUES (?, ?, ?, ?, ?, ?, ?, ?)
  154. RETURNING *
  155. """,
  156. [
  157. forum.topic,
  158. creator_id,
  159. forum.moderator_id,
  160. "pending",
  161. forum.duration_minutes,
  162. None,
  163. "[]",
  164. "{}"
  165. ]
  166. )
  167. db_forum = fetch_one(rs)
  168. tx.execute("DELETE FROM messages WHERE forum_id = ?", [db_forum.id])
  169. tx.execute("DELETE FROM forum_participants WHERE forum_id = ?", [db_forum.id])
  170. tx.execute("DELETE FROM system_logs WHERE forum_id = ?", [db_forum.id])
  171. if forum.participant_ids:
  172. unique_pids = list(dict.fromkeys(int(pid) for pid in forum.participant_ids))
  173. values = []
  174. placeholders = []
  175. for pid in unique_pids:
  176. placeholders.append("(?, ?, ?)")
  177. values.extend([db_forum.id, pid, "[]"])
  178. if values:
  179. query = f"INSERT INTO forum_participants (forum_id, persona_id, thoughts_history) VALUES {', '.join(placeholders)} ON CONFLICT (forum_id, persona_id) DO NOTHING"
  180. tx.execute(query, values)
  181. # FORCE COMMIT
  182. if hasattr(tx, 'commit'):
  183. tx.commit()
  184. elif hasattr(db, 'commit'):
  185. db.commit()
  186. # Return full object (will trigger cache set in get_forum)
  187. return get_forum(db, db_forum.id)
  188. except Exception as e:
  189. logger.error(f"Error creating forum: {e}")
  190. raise
  191. def delete_forum(db, forum_id: int):
  192. logger.info(f"Attempting to delete forum {forum_id}")
  193. try:
  194. with db_transaction(db) as tx:
  195. tx.execute("DELETE FROM messages WHERE forum_id = ?", [forum_id])
  196. tx.execute("DELETE FROM forum_participants WHERE forum_id = ?", [forum_id])
  197. tx.execute("DELETE FROM system_logs WHERE forum_id = ?", [forum_id])
  198. rs = tx.execute("DELETE FROM forums WHERE id = ?", [forum_id])
  199. affected = rs.rows_affected if hasattr(rs, 'rows_affected') else -1
  200. logger.info(f"Deleted forum {forum_id}, rows affected: {affected}")
  201. # FORCE COMMIT
  202. if hasattr(tx, 'commit'):
  203. tx.commit()
  204. logger.info("Transaction committed explicitly")
  205. elif hasattr(db, 'commit'):
  206. db.commit()
  207. logger.info("DB committed explicitly")
  208. success = affected > 0 if affected != -1 else True
  209. return success
  210. except Exception as e:
  211. logger.error(f"Error deleting forum: {e}")
  212. raise
  213. def get_forum(db, forum_id: int):
  214. rs = db.execute("SELECT * FROM forums WHERE id = ?", [forum_id])
  215. forum = fetch_one(rs)
  216. if not forum:
  217. return None
  218. participants = get_forum_participants(db, forum_id)
  219. setattr(forum, "participants", participants)
  220. if forum.moderator_id:
  221. mod_rs = db.execute("SELECT * FROM moderators WHERE id = ?", [forum.moderator_id])
  222. setattr(forum, "moderator", fetch_one(mod_rs))
  223. else:
  224. setattr(forum, "moderator", None)
  225. return forum
  226. def update_forum(
  227. db,
  228. forum_id: int,
  229. summary_history: list = None,
  230. status: str = None,
  231. start_time: datetime = None,
  232. ablation_flags: dict = None,
  233. ):
  234. try:
  235. set_clauses = []
  236. values = []
  237. if summary_history is not None:
  238. set_clauses.append("summary_history = ?")
  239. values.append(json.dumps(summary_history))
  240. if status is not None:
  241. set_clauses.append("status = ?")
  242. values.append(status)
  243. if start_time is not None:
  244. set_clauses.append("start_time = ?")
  245. values.append(start_time)
  246. if ablation_flags is not None:
  247. set_clauses.append("ablation_flags = ?")
  248. values.append(json.dumps(ablation_flags))
  249. if not set_clauses:
  250. return get_forum(db, forum_id)
  251. values.append(forum_id)
  252. query = f"UPDATE forums SET {', '.join(set_clauses)} WHERE id = ? RETURNING *"
  253. rs = db_execute_commit(db, query, values)
  254. updated = fetch_one(rs)
  255. return updated
  256. except Exception as e:
  257. logger.error(f"Error updating forum: {e}")
  258. raise
  259. def get_forum_participants(db, forum_id: int):
  260. query = """
  261. SELECT fp.*, p.name as persona_name, p.title as persona_title, p.bio as persona_bio,
  262. p.theories as persona_theories, p.stance as persona_stance,
  263. p.system_prompt as persona_system_prompt, p.owner_id as persona_owner_id,
  264. p.created_at as persona_created_at
  265. FROM forum_participants fp
  266. JOIN personas p ON fp.persona_id = p.id
  267. WHERE fp.forum_id = ?
  268. """
  269. rs = db.execute(query, [forum_id])
  270. rows = fetch_all(rs)
  271. results = []
  272. for row in rows:
  273. persona_data = {
  274. "id": row.persona_id,
  275. "name": row.persona_name,
  276. "title": row.persona_title,
  277. "bio": row.persona_bio,
  278. "theories": row.persona_theories,
  279. "stance": row.persona_stance,
  280. "system_prompt": row.persona_system_prompt,
  281. "owner_id": row.persona_owner_id,
  282. "created_at": row.persona_created_at
  283. }
  284. setattr(row, "persona", RowObject(persona_data))
  285. results.append(row)
  286. return results
  287. def update_forum_participant(db, forum_id: int, persona_id: int, thoughts_history: list = None):
  288. try:
  289. if thoughts_history is None:
  290. return None
  291. query = "UPDATE forum_participants SET thoughts_history = ? WHERE forum_id = ? AND persona_id = ? RETURNING *"
  292. rs = db_execute_commit(db, query, [json.dumps(thoughts_history), forum_id, persona_id])
  293. return fetch_one(rs)
  294. except Exception as e:
  295. logger.error(f"Error updating participant: {e}")
  296. raise
  297. def create_message(db, message: MessageCreate):
  298. try:
  299. timestamp = datetime.now()
  300. rs = db_execute_commit(
  301. db,
  302. """
  303. INSERT INTO messages (forum_id, persona_id, moderator_id, speaker_name, content, turn_count, thought, timestamp)
  304. VALUES (?, ?, ?, ?, ?, ?, ?, ?)
  305. RETURNING *
  306. """,
  307. [
  308. message.forum_id,
  309. message.persona_id,
  310. message.moderator_id,
  311. message.speaker_name,
  312. message.content,
  313. message.turn_count,
  314. message.thought,
  315. timestamp
  316. ]
  317. )
  318. return fetch_one(rs)
  319. except Exception as e:
  320. logger.error(f"Error creating message: {e}")
  321. raise
  322. def get_forum_messages(db, forum_id: int):
  323. rs = db.execute("SELECT * FROM messages WHERE forum_id = ? ORDER BY timestamp ASC", [forum_id])
  324. return fetch_all(rs)