mirror of
https://github.com/TandoorRecipes/recipes.git
synced 2026-01-01 04:10:06 -05:00
86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
from decimal import Decimal
|
|
|
|
from django.utils import timezone
|
|
from django.db.models import Sum
|
|
from litellm import CustomLogger
|
|
|
|
from cookbook.models import AiLog
|
|
from recipes import settings
|
|
|
|
|
|
def get_monthly_token_usage(space):
|
|
"""
|
|
returns the number of credits the space has used in the current month
|
|
"""
|
|
token_usage = AiLog.objects.filter(space=space, credits_from_balance=False, created_at__month=timezone.now().month).aggregate(Sum('credit_cost'))['credit_cost__sum']
|
|
if token_usage is None:
|
|
token_usage = 0
|
|
return token_usage
|
|
|
|
|
|
def has_monthly_token(space):
|
|
"""
|
|
checks if the monthly credit limit has been exceeded
|
|
"""
|
|
return get_monthly_token_usage(space) < space.ai_credits_monthly
|
|
|
|
|
|
def can_perform_ai_request(space):
|
|
return (has_monthly_token(space) or space.ai_credits_balance > 0) and space.ai_enabled
|
|
|
|
|
|
class AiCallbackHandler(CustomLogger):
|
|
space = None
|
|
user = None
|
|
ai_provider = None
|
|
function = None
|
|
|
|
def __init__(self, space, user, ai_provider, function):
|
|
super().__init__()
|
|
self.space = space
|
|
self.user = user
|
|
self.ai_provider = ai_provider
|
|
self.function = function
|
|
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
pass
|
|
|
|
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
|
|
pass
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
self.create_ai_log(kwargs, response_obj, start_time, end_time)
|
|
|
|
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
self.create_ai_log(kwargs, response_obj, start_time, end_time)
|
|
|
|
def create_ai_log(self, kwargs, response_obj, start_time, end_time):
|
|
credit_cost = 0
|
|
credits_from_balance = False
|
|
if self.ai_provider.log_credit_cost:
|
|
credit_cost = kwargs.get("response_cost", 0) * 100
|
|
|
|
if (not has_monthly_token(self.space)) and self.space.ai_credits_balance > 0:
|
|
remaining_balance = self.space.ai_credits_balance - Decimal(str(credit_cost))
|
|
if remaining_balance < 0:
|
|
remaining_balance = 0
|
|
if settings.HOSTED and self.space.ai_credits_monthly == 0:
|
|
self.space.ai_enabled = False
|
|
|
|
self.space.ai_credits_balance = remaining_balance
|
|
credits_from_balance = True
|
|
self.space.save()
|
|
|
|
AiLog.objects.create(
|
|
created_by=self.user,
|
|
space=self.space,
|
|
ai_provider=self.ai_provider,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
input_tokens=response_obj['usage']['prompt_tokens'],
|
|
output_tokens=response_obj['usage']['completion_tokens'],
|
|
function=self.function,
|
|
credit_cost=credit_cost,
|
|
credits_from_balance=credits_from_balance,
|
|
)
|