migration: preserve mentor and llm streaming slice

2026-07-31 04:18:53 +08:00
parent 4bab921d14
commit 2919229c73
26 changed files with 1705 additions and 1222 deletions
@@ -0,0 +1,102 @@
+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,
+            )