Files
Jarvis-Ai/src/javis/memory/usage_store.py
T

115 lines
4.2 KiB
Python

"""Metadata-only provider audit and local usage limits."""
from __future__ import annotations
import sqlite3
from contextlib import closing
from dataclasses import dataclass
from datetime import UTC, date, datetime
from pathlib import Path
class UsageStoreError(RuntimeError):
"""Provider metadata could not be stored or read."""
@dataclass(frozen=True, slots=True)
class ProviderEvent:
provider: str
model: str
success: bool
error_category: str | None
fallback: bool
privacy_policy: str
input_tokens: int | None = None
output_tokens: int | None = None
class SQLiteUsageStore:
def __init__(self, database_path: Path) -> None:
self.database_path = database_path
try:
self.database_path.parent.mkdir(parents=True, exist_ok=True)
with closing(self._connect()) as connection:
connection.execute(
"""
CREATE TABLE IF NOT EXISTS provider_events (
id INTEGER PRIMARY KEY AUTOINCREMENT,
occurred_at TEXT NOT NULL,
provider TEXT NOT NULL,
model TEXT NOT NULL,
success INTEGER NOT NULL,
error_category TEXT,
fallback INTEGER NOT NULL,
privacy_policy TEXT NOT NULL,
input_tokens INTEGER,
output_tokens INTEGER
)
"""
)
connection.execute(
"""
CREATE INDEX IF NOT EXISTS idx_provider_events_time_provider
ON provider_events(occurred_at, provider)
"""
)
connection.commit()
except (OSError, sqlite3.Error) as exc:
raise UsageStoreError(
f"Die Nutzungsmetadatenbank kann nicht geöffnet werden: {database_path}"
) from exc
def _connect(self) -> sqlite3.Connection:
return sqlite3.connect(self.database_path, timeout=10)
def record(self, event: ProviderEvent) -> None:
try:
with closing(self._connect()) as connection:
connection.execute(
"""
INSERT INTO provider_events (
occurred_at, provider, model, success, error_category,
fallback, privacy_policy, input_tokens, output_tokens
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
datetime.now(UTC).isoformat(timespec="seconds"),
event.provider,
event.model,
int(event.success),
event.error_category,
int(event.fallback),
event.privacy_policy,
event.input_tokens,
event.output_tokens,
),
)
connection.commit()
except sqlite3.Error as exc:
raise UsageStoreError("Nutzungsmetadaten konnten nicht gespeichert werden.") from exc
def cloud_requests_on(self, day: date | None = None) -> int:
return self._count(day or datetime.now(UTC).date(), "provider = 'gemini'")
def local_fallbacks_on(self, day: date | None = None) -> int:
return self._count(
day or datetime.now(UTC).date(),
"provider = 'ollama' AND fallback = 1",
)
def _count(self, day: date, condition: str) -> int:
start = f"{day.isoformat()}T00:00:00+00:00"
end = f"{day.isoformat()}T23:59:59+00:00"
try:
with closing(self._connect()) as connection:
row = connection.execute(
f"""
SELECT COUNT(*) FROM provider_events
WHERE occurred_at BETWEEN ? AND ? AND {condition}
""",
(start, end),
).fetchone()
except sqlite3.Error as exc:
raise UsageStoreError("Nutzungsmetadaten konnten nicht gelesen werden.") from exc
return int(row[0])