[Prototype] Daily/Weekly statistics for Grafana Infinity plugin

README.md

@@ -98,6 +98,7 @@ metrics:
   environment:
     - MONGODB_URI=mongodb://mongodb:27017/
     - LOGGING_LEVEL=info
+    - TZ=Europe/Berlin
   restart: unless-stopped
 ```
 

metrics.py

@@ -1,14 +1,20 @@
-from pymongo import MongoClient
-from prometheus_client import start_http_server
-from prometheus_client.core import GaugeMetricFamily, REGISTRY
-from prometheus_client.registry import Collector
+import json
 import logging
 import os
-import signal
-import sys
+import shutil
+import tempfile
 import time
-from datetime import datetime, timezone, timedelta
+from datetime import datetime, timedelta, timezone
+
 from dotenv import load_dotenv
+from prometheus_client import start_http_server
+from prometheus_client.core import REGISTRY, GaugeMetricFamily
+from prometheus_client.registry import Collector
+from prometheus_client.twisted import MetricsResource
+from pymongo import MongoClient
+from twisted.internet import reactor
+from twisted.web.resource import Resource
+from twisted.web.server import Site
 
 load_dotenv()
 
@@ -23,12 +29,11 @@ class LibreChatMetricsCollector(Collector):
     A custom Prometheus collector that gathers metrics from the LibreChat MongoDB database.
     """
 
-    def __init__(self, mongodb_uri):
+    def __init__(self, mongodb):
         """
         Initialize the MongoDB client and set up initial state.
         """
-        self.client = MongoClient(mongodb_uri)
-        self.db = self.client["LibreChat"]
+        self.db = mongodb
         self.messages_collection = self.db["messages"]
 
     def collect(self):
@@ -334,33 +339,130 @@ def collect_uploaded_file_count(self):
             logger.exception("Error collecting uploaded files: %s", e)
 
 
-def signal_handler(sig, frame):
-    """
-    Handle termination signals to allow for graceful shutdown.
-    """
-    logger.info("Shutting down gracefully...")
-    sys.exit(0)
+class InfinityResource(Resource):
+    isLeaf = True
+    stat_types = ("daily", "weekly", "monthly", "yearly")
+
+    def __init__(self, mongodb):
+        """
+        Initialize the MongoDB client and set up initial state.
+        """
+        self.messages = mongodb["messages"]
+
+        # Load stored statistics
+        try:
+            statistics_file = os.getenv("STATISTICS_FILE", "statistics.json")
+            with open(statistics_file, "rb") as f:
+                logger.info("Loading cached statistics from %s", statistics_file)
+                self.statistics = json.load(f)
+        except FileNotFoundError:
+            logger.debug("No cached statistics exist")
+            self.statistics = {t: [] for t in self.stat_types}
+
+        # Initial update of statistics
+        self.update_statistics()
+
+        self.save_cache()
+
+    def save_cache(self):
+        statistics_file = os.getenv("STATISTICS_FILE", "statistics.json")
+        logger.info("Saving statistics to %s", statistics_file)
+        with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp:
+            json.dump(self.statistics, temp)
+            temp_file = temp.name
+        logger.debug("Moving %s to %s", temp_file, statistics_file)
+        shutil.move(temp_file, statistics_file)
+
+    def next_step(self, date, stat_type):
+        if stat_type == "yearly":
+            return datetime(year=date.year + 1, month=1, day=1)
+        elif stat_type == "monthly":
+            date = date.replace(day=15) + timedelta(days=30)
+            return datetime(year=date.year, month=date.month, day=1)
+        elif stat_type == "weekly":
+            date = date + timedelta(days=7 - date.weekday())
+            return datetime(year=date.year, month=date.month, day=date.day)
+        elif stat_type == "daily":
+            date = date + timedelta(days=1)
+            return datetime(year=date.year, month=date.month, day=date.day)
+        raise Exception(f"Unknown stat type: {stat_type}")
+
+    def statistics_start(self, stat_type):
+        stats = self.statistics.get(stat_type)
+        if stats:
+            latest = stats[-1].get("date")
+            date = datetime.fromisoformat(latest)
+            return self.next_step(date, stat_type)
+        else:
+            self.statistics[stat_type] = []
+            # Get oldest message to use this as a start point for getting daily statistics
+            start = datetime.today()
+            for message in self.messages.find().sort("createdAt", 1).limit(1):
+                start = message["createdAt"]
+            if stat_type == "yearly":
+                return datetime(year=start.year, month=1, day=1)
+            elif stat_type == "monthly":
+                return datetime(year=start.year, month=start.month, day=1)
+            elif stat_type == "weekly":
+                start = start - timedelta(days=start.weekday())
+                return datetime(year=start.year, month=start.month, day=start.day)
+            return datetime(year=start.year, month=start.month, day=start.day)
+
+    def update_statistics(self):
+        for stat_type in self.stat_types:
+            today = datetime.today()
+            start = self.statistics_start(stat_type)
+            end = self.next_step(start, stat_type)
+            stats = self.statistics[stat_type]
+            while end < today:
+                logger.info("Getting %s users from %s", stat_type, start)
+                count = self.user_count(start, end)
+                stats.append({"date": start.date().isoformat(), "user": count})
+                start = end
+                end = self.next_step(end, stat_type)
+        logger.debug("Statistics %s", self.statistics)
+
+    def user_count(self, start, end):
+        return len(
+            self.messages.distinct("user", {"createdAt": {"$gte": start, "$lt": end}})
+        )
+
+    def render_GET(self, request):
+        # Update statistics if necessary
+        self.update_statistics()
+
+        # Get time frame parameters
+        begin_ts = float(request.args.get(b"from", [0])[0]) / 1000.0
+        begin = datetime.fromtimestamp(begin_ts)
+        end_ts = float(request.args.get(b"to", [time.time() * 1000.0])[0]) / 1000.0
+        end = datetime.fromtimestamp(end_ts)
+        logger.info("from: %s -> %s", begin, end)
+
+        # Set the response code and content type
+        request.setHeader(b"Content-Type", b"application/json")
+        request.setResponseCode(200)  # HTTP status code
+
+        # Return the response as JSON
+        return json.dumps(self.statistics).encode("utf-8")
 
 
 if __name__ == "__main__":
-    # Get MongoDB URI and Prometheus port from environment variables
+    # Get MongoDB URI and initialize database connection
     mongodb_uri = os.getenv("MONGODB_URI", "mongodb://mongodb:27017/")
-    prometheus_port = 8000
+    mongodb = MongoClient(mongodb_uri)["LibreChat"]
 
-    # Handle shutdown signals
-    signal.signal(signal.SIGINT, signal_handler)
-    signal.signal(signal.SIGTERM, signal_handler)
+    port = 8000
 
     # Start the Prometheus exporter
-    collector = LibreChatMetricsCollector(mongodb_uri)
+    collector = LibreChatMetricsCollector(mongodb)
     REGISTRY.register(collector)
-    start_http_server(prometheus_port)
-    logger.info(f"Metrics server is running on port {prometheus_port}.")
-
-    # Keep the application running
-    try:
-        while True:
-            time.sleep(1)
-    except KeyboardInterrupt:
-        logger.info("Received interrupt, shutting down.")
-        sys.exit(0)
+    logger.info("Starting server on port %i", port)
+
+    root = Resource()
+    metrics = MetricsResource()
+    root.putChild(b"", metrics)
+    root.putChild(b"metrics", metrics)
+    root.putChild(b"infinity", InfinityResource(mongodb))
+
+    reactor.listenTCP(port, Site(root))
+    reactor.run()

requirements.txt

@@ -1,3 +1,4 @@
 prometheus-client
 pymongo
 python-dotenv
+Twisted