136 lines
4.2 KiB
Python
136 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
|
|
|
|
class LLMRepository:
|
|
@staticmethod
|
|
def usage_today(connection: sqlite3.Connection, user_id: int, usage_date: str) -> int:
|
|
row = connection.execute(
|
|
"SELECT successful_calls FROM llm_usage_daily WHERE user_id = ? AND usage_date = ?",
|
|
(user_id, usage_date),
|
|
).fetchone()
|
|
return int(row["successful_calls"]) if row else 0
|
|
|
|
@staticmethod
|
|
def usage_for_users(connection: sqlite3.Connection, usage_date: str) -> dict[int, int]:
|
|
return {
|
|
int(row["user_id"]): int(row["successful_calls"])
|
|
for row in connection.execute(
|
|
"SELECT user_id, successful_calls FROM llm_usage_daily WHERE usage_date = ?",
|
|
(usage_date,),
|
|
)
|
|
}
|
|
|
|
@staticmethod
|
|
def active_today(connection: sqlite3.Connection, user_id: int, day_prefix: str) -> int:
|
|
row = connection.execute(
|
|
"""
|
|
SELECT COUNT(*) AS total FROM llm_requests
|
|
WHERE user_id = ? AND started_at LIKE ? AND status IN ('reserved', 'streaming')
|
|
""",
|
|
(user_id, f"{day_prefix}%"),
|
|
).fetchone()
|
|
return int(row["total"])
|
|
|
|
@staticmethod
|
|
def reserve(
|
|
connection: sqlite3.Connection,
|
|
*,
|
|
request_id: str,
|
|
user_id: int,
|
|
feature: str,
|
|
business_id: str,
|
|
prompt_version: str,
|
|
started_at: str,
|
|
input_chars: int,
|
|
) -> None:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO llm_requests (
|
|
id, user_id, feature, business_id, prompt_version,
|
|
status, started_at, input_chars
|
|
) VALUES (?, ?, ?, ?, ?, 'reserved', ?, ?)
|
|
""",
|
|
(
|
|
request_id,
|
|
user_id,
|
|
feature,
|
|
business_id,
|
|
prompt_version,
|
|
started_at,
|
|
input_chars,
|
|
),
|
|
)
|
|
|
|
@staticmethod
|
|
def set_request_status(
|
|
connection: sqlite3.Connection,
|
|
request_id: str,
|
|
status: str,
|
|
*,
|
|
completed_at: str | None = None,
|
|
duration_ms: int = 0,
|
|
error_type: str = "",
|
|
output_chars: int = 0,
|
|
) -> None:
|
|
connection.execute(
|
|
"""
|
|
UPDATE llm_requests SET status = ?, completed_at = ?, duration_ms = ?,
|
|
error_type = ?, output_chars = ? WHERE id = ?
|
|
""",
|
|
(status, completed_at, duration_ms, error_type, output_chars, request_id),
|
|
)
|
|
|
|
@staticmethod
|
|
def increment_usage(
|
|
connection: sqlite3.Connection, user_id: int, usage_date: str, updated_at: str
|
|
) -> None:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO llm_usage_daily (user_id, usage_date, successful_calls, updated_at)
|
|
VALUES (?, ?, 1, ?)
|
|
ON CONFLICT(user_id, usage_date) DO UPDATE SET
|
|
successful_calls = successful_calls + 1,
|
|
updated_at = excluded.updated_at
|
|
""",
|
|
(user_id, usage_date, updated_at),
|
|
)
|
|
|
|
@staticmethod
|
|
def start_attempt(
|
|
connection: sqlite3.Connection,
|
|
request_id: str,
|
|
model_id: int,
|
|
role: str,
|
|
started_at: str,
|
|
input_chars: int,
|
|
) -> int:
|
|
cursor = connection.execute(
|
|
"""
|
|
INSERT INTO llm_attempts (
|
|
request_id, model_id, role, status, started_at, input_chars
|
|
) VALUES (?, ?, ?, 'streaming', ?, ?)
|
|
""",
|
|
(request_id, model_id, role, started_at, input_chars),
|
|
)
|
|
return int(cursor.lastrowid)
|
|
|
|
@staticmethod
|
|
def finish_attempt(
|
|
connection: sqlite3.Connection,
|
|
attempt_id: int,
|
|
status: str,
|
|
completed_at: str,
|
|
duration_ms: int,
|
|
error_type: str,
|
|
output_chars: int,
|
|
) -> None:
|
|
connection.execute(
|
|
"""
|
|
UPDATE llm_attempts SET status = ?, completed_at = ?, duration_ms = ?,
|
|
error_type = ?, output_chars = ? WHERE id = ?
|
|
""",
|
|
(status, completed_at, duration_ms, error_type, output_chars, attempt_id),
|
|
)
|