|
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Any
|
|
|
|
from devplacepy.database import db, get_table
|
|
|
|
logger = logging.getLogger("devii.store")
|
|
|
|
CONVERSATIONS = "devii_conversations"
|
|
LEDGER = "devii_usage_ledger"
|
|
TURNS = "devii_turns"
|
|
|
|
PROMPT_AUDIT_CHARS = 4000
|
|
REPLY_AUDIT_CHARS = 4000
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def _iso(moment: datetime) -> str:
|
|
return moment.isoformat()
|
|
|
|
|
|
class ConversationStore:
|
|
def load(self, owner_kind: str, owner_id: str) -> list[dict[str, Any]] | None:
|
|
if CONVERSATIONS not in db.tables:
|
|
return None
|
|
row = get_table(CONVERSATIONS).find_one(owner_kind=owner_kind, owner_id=owner_id)
|
|
if not row or not row.get("messages"):
|
|
return None
|
|
try:
|
|
return json.loads(row["messages"])
|
|
except (ValueError, TypeError):
|
|
logger.warning("Corrupt conversation for %s/%s", owner_kind, owner_id)
|
|
return None
|
|
|
|
def save(self, owner_kind: str, owner_id: str, messages: list[dict[str, Any]]) -> None:
|
|
now = _iso(_now())
|
|
record = {
|
|
"owner_kind": owner_kind,
|
|
"owner_id": owner_id,
|
|
"messages": json.dumps(messages),
|
|
"updated_at": now,
|
|
}
|
|
table = get_table(CONVERSATIONS)
|
|
existing = table.find_one(owner_kind=owner_kind, owner_id=owner_id) if CONVERSATIONS in db.tables else None
|
|
if existing:
|
|
table.update({**record, "id": existing["id"]}, ["id"])
|
|
else:
|
|
record["created_at"] = now
|
|
table.insert(record)
|
|
|
|
def clear(self, owner_kind: str, owner_id: str) -> None:
|
|
if CONVERSATIONS in db.tables:
|
|
get_table(CONVERSATIONS).delete(owner_kind=owner_kind, owner_id=owner_id)
|
|
|
|
|
|
class UsageLedger:
|
|
def record(
|
|
self,
|
|
owner_kind: str,
|
|
owner_id: str,
|
|
turn_id: str,
|
|
usage: dict[str, int],
|
|
cost_usd: float,
|
|
model: str,
|
|
) -> None:
|
|
get_table(LEDGER).insert({
|
|
"owner_kind": owner_kind,
|
|
"owner_id": owner_id,
|
|
"turn_id": turn_id,
|
|
"created_at": _iso(_now()),
|
|
"prompt_tokens": int(usage.get("prompt_tokens", 0)),
|
|
"completion_tokens": int(usage.get("completion_tokens", 0)),
|
|
"cache_hit_tokens": int(usage.get("cache_hit_tokens", 0)),
|
|
"cache_miss_tokens": int(usage.get("cache_miss_tokens", 0)),
|
|
"cost_usd": float(cost_usd),
|
|
"model": model,
|
|
})
|
|
|
|
def spent_24h(self, owner_kind: str, owner_id: str) -> float:
|
|
if LEDGER not in db.tables:
|
|
return 0.0
|
|
cutoff = _iso(_now() - timedelta(hours=24))
|
|
rows = db.query(
|
|
f"SELECT COALESCE(SUM(cost_usd), 0) AS total FROM {LEDGER} "
|
|
"WHERE owner_kind = :owner_kind AND owner_id = :owner_id AND created_at >= :cutoff",
|
|
owner_kind=owner_kind, owner_id=owner_id, cutoff=cutoff,
|
|
)
|
|
for row in rows:
|
|
return float(row["total"] or 0.0)
|
|
return 0.0
|
|
|
|
def turns_24h(self, owner_kind: str, owner_id: str) -> int:
|
|
if LEDGER not in db.tables:
|
|
return 0
|
|
cutoff = _iso(_now() - timedelta(hours=24))
|
|
rows = db.query(
|
|
f"SELECT COUNT(*) AS n FROM {LEDGER} "
|
|
"WHERE owner_kind = :owner_kind AND owner_id = :owner_id AND created_at >= :cutoff",
|
|
owner_kind=owner_kind, owner_id=owner_id, cutoff=cutoff,
|
|
)
|
|
for row in rows:
|
|
return int(row["n"] or 0)
|
|
return 0
|
|
|
|
def total_spent_24h(self) -> float:
|
|
if LEDGER not in db.tables:
|
|
return 0.0
|
|
cutoff = _iso(_now() - timedelta(hours=24))
|
|
rows = db.query(
|
|
f"SELECT COALESCE(SUM(cost_usd), 0) AS total FROM {LEDGER} WHERE created_at >= :cutoff",
|
|
cutoff=cutoff,
|
|
)
|
|
for row in rows:
|
|
return float(row["total"] or 0.0)
|
|
return 0.0
|
|
|
|
def prune(self, older_than_hours: int = 48) -> int:
|
|
if LEDGER not in db.tables:
|
|
return 0
|
|
cutoff = _iso(_now() - timedelta(hours=older_than_hours))
|
|
return int(get_table(LEDGER).delete(created_at={"<": cutoff}))
|
|
|
|
def reset(self, owner_kind: str, owner_id: str) -> int:
|
|
if LEDGER not in db.tables:
|
|
return 0
|
|
table = get_table(LEDGER)
|
|
count = table.count(owner_kind=owner_kind, owner_id=owner_id)
|
|
table.delete(owner_kind=owner_kind, owner_id=owner_id)
|
|
return int(count)
|
|
|
|
def reset_owner_kind(self, owner_kind: str) -> int:
|
|
if LEDGER not in db.tables:
|
|
return 0
|
|
table = get_table(LEDGER)
|
|
count = table.count(owner_kind=owner_kind)
|
|
table.delete(owner_kind=owner_kind)
|
|
return int(count)
|
|
|
|
def reset_all(self) -> int:
|
|
if LEDGER not in db.tables:
|
|
return 0
|
|
table = get_table(LEDGER)
|
|
count = table.count()
|
|
table.delete()
|
|
return int(count)
|
|
|
|
|
|
class TurnAudit:
|
|
def record(self, record: dict[str, Any]) -> None:
|
|
record = dict(record)
|
|
record["prompt"] = (record.get("prompt") or "")[:PROMPT_AUDIT_CHARS]
|
|
record["reply"] = (record.get("reply") or "")[:REPLY_AUDIT_CHARS]
|
|
get_table(TURNS).insert(record)
|