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,46 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+
+class LLMAuditRepositoryMixin:
+    def record_llm_usage(
+        self,
+        user_id: int,
+        feature: str,
+        source: str,
+        model: str,
+        status: str,
+        latency_ms: int = 0,
+        *,
+        role: str = "",
+        prompt_version: str = "",
+        error_code: str = "",
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+    ) -> None:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        with self.connect() as connection:
+            connection.execute(
+                """
+                INSERT INTO llm_usage
+                    (user_id, feature, source, model, status, latency_ms, created_at,
+                     role, prompt_version, error_code, input_tokens, output_tokens)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    user_id, feature, source, model, status, int(latency_ms), now,
+                    role, prompt_version, error_code, int(input_tokens), int(output_tokens),
+                ),
+            )
+
+    def count_llm_usage_since(self, user_id: int, source: str, since: str) -> int:
+        with self.connect() as connection:
+            row = connection.execute(
+                """
+                SELECT COUNT(*) AS total FROM llm_usage
+                WHERE user_id = ? AND source = ? AND created_at >= ?
+                """,
+                (user_id, source, since),
+            ).fetchone()
+        return int(row["total"] if row else 0)