add time estimation for image gen

vibe_bot/database.py

@@ -30,6 +30,11 @@ logging.basicConfig(
 )
 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
+
 
 class ChatDatabase:
     """SQLite database with RAG support for storing chat history
@@ -118,6 +123,19 @@ class ChatDatabase:
         )
         logger.info("idx_user_id index created successfully")
 
+        # Create image generation timing table
+        logger.info("Creating image_generation_times table if not exists")
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS image_generation_times (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                duration_seconds REAL NOT NULL,
+                generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        """,
+        )
+        logger.info("image_generation_times table initialized successfully")
+
         conn.commit()
         logger.info("Database initialization completed successfully")
         conn.close()
@@ -497,6 +515,85 @@ class ChatDatabase:
         conn.commit()
         conn.close()
 
+    def record_image_generation_time(self, duration_seconds: float) -> bool:
+        """Record how long an image generation took.
+
+        Args:
+            duration_seconds: Wall-clock seconds the generation took.
+
+        """
+        logger.info("Recording image generation time: %.2fs", duration_seconds)
+        conn = sqlite3.connect(self.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(self) -> 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 = sqlite3.connect(self.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])
+
 
 # Global database instance
 _chat_db: ChatDatabase | None = None