| 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() | |