47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
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)
|