ai4data-mcp / utils.py
avsolatorio's picture
Add logging
34b2550
import requests
import json
import os
import torch
import requests
import json
from datetime import datetime, timezone
from dotenv import load_dotenv
load_dotenv()
WEBHOOK_URL = os.getenv("WEBHOOK_URL")
WEBHOOK_SECRET = os.getenv("WEBHOOK_SECRET")
def get_best_torch_device():
if torch.cuda.is_available():
return torch.device("cuda")
elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
return torch.device("mps")
else:
return torch.device("cpu")
device = get_best_torch_device()
def send_post(payload: dict):
"""
Send a post request to the webhook.
"""
headers = {
"Content-Type": "application/json",
"X-Log-Secret": WEBHOOK_SECRET,
}
resp = requests.post(WEBHOOK_URL, headers=headers, data=json.dumps(payload))
if resp.status_code == 200:
return True
else:
print(f"Posting to webhook failed: {resp.status_code} {resp.text}")
return False
def hf_send_post(payload: dict):
"""
Send a post request to the HF webhook.
"""
payload["service"] = "hf-ai4data-mcp-server"
payload["level"] = "INFO"
payload["timestamp"] = datetime.now(timezone.utc).isoformat()
return send_post(payload)
# Example usage
if __name__ == "__main__":
from datetime import datetime
post_entry = {
"service": "hf-ai4data-mcp-api-test",
"level": "INFO",
"message": "Test message " + datetime.now().isoformat(),
}
send_post(post_entry)