axiom logger

axiom_logger.py

@@ -0,0 +1,66 @@
+import axiom
+from datetime import datetime,timedelta
+import rfc3339
+from dotenv import load_dotenv
+import os
+
+
+
+class AxiomLogger:
+
+    def __init__(self):
+        load_dotenv()
+        api_token = os.getenv('AXIOM_API_KEY')
+        org_id = os.getenv('AXIOM_ORG_ID')
+        self.client = axiom.Client(api_token, org_id)
+        self.dataset='sdxl-server'
+
+    def get_timestamp_formatted(self):
+        time = datetime.utcnow()
+        time_formatted = rfc3339.format(time)
+        return time_formatted
+    
+    def ingest_with_status(self, status, message, **kwargs):
+        self.client.ingest_events(
+            dataset=self.dataset,
+            events=[{
+                "_time": self.get_timestamp_formatted(),
+                "status": status,
+                "message": message,
+                **kwargs
+            }])
+
+    def warning(self, message, **kwargs):
+        self.ingest_with_status("WARNING", message, **kwargs)
+    
+    def error(self, message, **kwargs):
+        self.ingest_with_status("ERROR", message, **kwargs)
+    
+    def info(self, message, **kwargs):
+        self.ingest_with_status("INFO", message, **kwargs)
+
+    def debug(self, message, **kwargs):
+        self.ingest_with_status("DEBUG", message, **kwargs)
+        
+    def critical(self, message, **kwargs):
+        self.ingest_with_status("CRITICAL", message, **kwargs)
+
+
+def main():
+    api_token = os.getenv('AXIOM_API_KEY')
+    org_id = os.getenv('AXIOM_ORG_ID')
+    client = axiom.Client(api_token, org_id)
+
+    time = datetime.utcnow() - timedelta(hours=1)
+    time_formatted = rfc3339.format(time)
+
+    client.ingest_events(
+        dataset="sdxl-server",
+        events=[
+            {"inference_test": "test1", "_time": time_formatted},
+            {"inference_test": "test2", "_time": time_formatted},
+        ])
+
+
+if __name__=='__main__':
+    main()