climateqa/logging.py

import datetime
import json
import os

from azure.storage.fileshare import ShareServiceClient


def log(question, history, docs, user_id):
    if has_blob_config():
        log_in_azure(question, history, docs, user_id)
    pass


def has_blob_config():
    """
    Checks if the necessary environment variables for Azure Blob Storage are set.
    Returns True if they are set, False otherwise.
    """
    return all(
        key in os.environ
        for key in ["BLOB_ACCOUNT_KEY", "BLOB_ACCOUNT_NAME", "BLOB_ACCOUNT_URL"]
    )


def log_in_azure(question, history, docs, user_id):
    timestamp = str(datetime.now().timestamp())
    file_name = timestamp + ".json"
    prompt = history[-1][0]
    logs = {
        "user_id": str(user_id),
        "prompt": prompt,
        "query": prompt,
        "question": question,
        "docs": docs,
        "answer": history[-1][1],
        "time": timestamp,
    }
    upload_azure(file_name, logs)


def get_azure_blob_client():
    account_key = os.environ["BLOB_ACCOUNT_KEY"]
    if len(account_key) == 86:
        account_key += "=="

    credential = {
        "account_key": account_key,
        "account_name": os.environ["BLOB_ACCOUNT_NAME"],
    }
    account_url = os.environ["BLOB_ACCOUNT_URL"]
    file_share_name = "climategpt"
    service = ShareServiceClient(account_url=account_url, credential=credential)
    share_client = service.get_share_client(file_share_name)
    return share_client

if has_blob_config():
    share_client = get_azure_blob_client()


def upload_azure(file, logs):
    logs = json.dumps(logs)
    print(type(logs))
    assert share_client is not None
    file_client = share_client.get_file_client(file)
    print("Uploading logs to Azure Blob Storage")
    print("----------------------------------")
    print("")
    print(logs)
    file_client.upload_file(logs)
    print("Logs uploaded to Azure Blob Storage")