from __future__ import annotations from datetime import datetime from typing import Any class MentorRepositoryMixin: def save_mentor_exchange( self, user_id: int, mentor_id: str, trade_date: str, question: str, answer: str, meta: str = "", ) -> None: now = datetime.now().astimezone().isoformat(timespec="seconds") with self.connect() as connection: connection.executemany( """ INSERT INTO mentor_messages (user_id, mentor_id, trade_date, role, content, meta, created_at) VALUES (?, ?, ?, ?, ?, ?, ?) """, [ (int(user_id), mentor_id, trade_date, "user", question, "", now), (int(user_id), mentor_id, trade_date, "assistant", answer, meta, now), ], ) connection.execute( """ DELETE FROM mentor_messages WHERE user_id = ? AND id NOT IN ( SELECT id FROM mentor_messages WHERE user_id = ? ORDER BY id DESC LIMIT 500 ) """, (int(user_id), int(user_id)), ) def list_mentor_messages( self, user_id: int, mentor_id: str, trade_date: str, limit: int = 100 ) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute( """ SELECT role, content, meta, created_at FROM mentor_messages WHERE user_id = ? AND mentor_id = ? AND trade_date = ? ORDER BY id DESC LIMIT ? """, (int(user_id), mentor_id, trade_date, max(1, min(500, int(limit)))), ).fetchall() return [dict(row) for row in reversed(rows)] def delete_mentor_messages(self, user_id: int, mentor_id: str, trade_date: str) -> int: with self.connect() as connection: cursor = connection.execute( "DELETE FROM mentor_messages WHERE user_id = ? AND mentor_id = ? AND trade_date = ?", (int(user_id), mentor_id, trade_date), ) return int(cursor.rowcount) def list_mentor_preferences(self, user_id: int) -> list[dict[str, Any]]: with self.connect() as connection: rows = connection.execute( """ SELECT mentor_id, pinned, sort_order FROM mentor_preferences WHERE user_id = ? ORDER BY sort_order, mentor_id """, (int(user_id),), ).fetchall() return [ { "mentor_id": str(row["mentor_id"]), "pinned": bool(row["pinned"]), "sort_order": int(row["sort_order"]), } for row in rows ] def save_mentor_preferences( self, user_id: int, ordered_ids: list[str], pinned_ids: set[str] ) -> None: now = datetime.now().astimezone().isoformat(timespec="seconds") values = [ (int(user_id), mentor_id, int(mentor_id in pinned_ids), index, now) for index, mentor_id in enumerate(ordered_ids) ] with self.connect() as connection: connection.execute( "DELETE FROM mentor_preferences WHERE user_id = ?", (int(user_id),), ) connection.executemany( """ INSERT INTO mentor_preferences (user_id, mentor_id, pinned, sort_order, updated_at) VALUES (?, ?, ?, ?, ?) """, values, )