WIP AI system

cookbook/helper/ai_helper.py

@@ -1,5 +1,6 @@
 from django.utils import timezone
 from django.db.models import Sum
+from litellm import CustomLogger
 
 from cookbook.models import AiLog
 
@@ -22,4 +23,60 @@ def has_monthly_token(space):
 
 
 def can_perform_ai_request(space):
-    return has_monthly_token(space) and space.ai_enabled
+    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
+
+    def __init__(self, space, user, ai_provider):
+        super().__init__()
+        self.space = space
+        self.user = user
+        self.ai_provider = ai_provider
+
+    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
+
+        print(not has_monthly_token(self.space) , self.space.ai_credits_balance > 0, not has_monthly_token(self.space) and self.space.ai_credits_balance > 0)
+        if (not has_monthly_token(self.space)) and self.space.ai_credits_balance > 0:
+            print('taking credits from balance')
+            self.space.ai_credits_balance = max(0, self.space.ai_credits_balance - credit_cost)
+            print('setting from balance to true')
+            credits_from_balance = True
+            print('saving space')
+            self.space.save()
+            print('done')
+        else:
+            print('not taking credits from balance')
+
+        print('creating AI log with credit cost ', credit_cost , ' from balance: ', credits_from_balance)
+        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=AiLog.F_FILE_IMPORT,
+            credit_cost=credit_cost,
+            credits_from_balance=credits_from_balance,
+        )