103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
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,
|
|
)
|