96 lines
2.5 KiB
Python
96 lines
2.5 KiB
Python
"""Image-generation timing statistics (moving-average estimate)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
from vibe_bot.db.connection import connect
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Moving-average window (most recent generations) for the time estimate.
|
|
IMAGE_GEN_TIME_WINDOW = 10
|
|
# Maximum number of generation times retained in the database.
|
|
IMAGE_GEN_TIME_LIMIT = 100
|
|
|
|
|
|
def record_image_generation_time(db_path: str, duration_seconds: float) -> bool:
|
|
"""Record how long an image generation took.
|
|
|
|
Args:
|
|
duration_seconds: Wall-clock seconds the generation took.
|
|
|
|
"""
|
|
logger.debug("Recording image generation time: %.2fs", duration_seconds)
|
|
conn = connect(db_path)
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO image_generation_times (duration_seconds)
|
|
VALUES (?)
|
|
""",
|
|
(duration_seconds,),
|
|
)
|
|
|
|
# Cap the table so it doesn't grow unbounded.
|
|
cursor.execute(
|
|
"""
|
|
DELETE FROM image_generation_times
|
|
WHERE id NOT IN (
|
|
SELECT id FROM image_generation_times
|
|
ORDER BY id DESC
|
|
LIMIT ?
|
|
)
|
|
""",
|
|
(IMAGE_GEN_TIME_LIMIT,),
|
|
)
|
|
|
|
conn.commit()
|
|
except Exception:
|
|
logger.exception("Error recording image generation time")
|
|
conn.rollback()
|
|
return False
|
|
else:
|
|
return True
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def get_image_generation_time_estimate(db_path: str) -> float | None:
|
|
"""Get a moving-average estimate of image generation time.
|
|
|
|
Returns:
|
|
The average duration in seconds over the most recent generations,
|
|
or None if there is no generation history yet.
|
|
|
|
"""
|
|
logger.debug("Computing moving-average image generation time estimate")
|
|
conn = connect(db_path)
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
cursor.execute(
|
|
"""
|
|
SELECT AVG(duration_seconds)
|
|
FROM (
|
|
SELECT duration_seconds
|
|
FROM image_generation_times
|
|
ORDER BY id DESC
|
|
LIMIT ?
|
|
)
|
|
""",
|
|
(IMAGE_GEN_TIME_WINDOW,),
|
|
)
|
|
row = cursor.fetchone()
|
|
except Exception:
|
|
logger.exception("Error reading image generation times")
|
|
return None
|
|
finally:
|
|
conn.close()
|
|
|
|
if row is None or row[0] is None:
|
|
return None
|
|
return float(row[0])
|