Sweep: Make a counter system to AI calls in upsonic_on_prem/utils/ai/ai.py (✓ Sandbox Passed)

upsonic_on_prem/utils/ai/ai.py

@@ -18,6 +18,7 @@
 from langchain_community.vectorstores import Chroma
 
 from upsonic_on_prem.utils import storage
+from upsonic_on_prem.utils.scope import storage_4
 
 from upsonic_on_prem.utils import debug, info, warning, failed, successfully
 
@@ -123,6 +124,19 @@
 
     def completion(self, input_text, model):
         result = None
+        import json
+        import traceback
+        try:
+            self.model_call_counter[model] = self.model_call_counter.get(model, 0) + 1
+            call_count = self.model_call_counter[model]
+            if call_count % 10 == 0:
+                debug(f"Model {model} has been called {call_count} times.")
+
+            storage_4.set("ai_model_call_counts", json.dumps(self.model_call_counter))
+        except Exception as e:
+            debug(f"Failed to persist model call counts: {str(e)}")
+        if call_count % 10 == 0:
+            debug(f"Model {model} has been called {call_count} times.")
         if model == "gemma-2b":
             result = self.gemmma(input_text)
         elif model == "gpt-3.5-turbo":
@@ -361,4 +375,6 @@
 
 
 
+    model_call_counter = {}
+
 AI = AI_()

upsonic_on_prem/utils/db/__init__.py

@@ -35,6 +35,7 @@ def can_access_to_config():
         storage = redis_client_()
         storage_2 = redis_client_(db=1)
         storage_3 = redis_client_(db=2)
+        storage_4 = redis_client_(db=3)
         succed = True
     except:
         warning("Redis configuration failed retrying in 5 seconds")