diff --git a/upsonic_on_prem/utils/ai/ai.py b/upsonic_on_prem/utils/ai/ai.py index df24c2df..f38ed77e 100644 --- a/upsonic_on_prem/utils/ai/ai.py +++ b/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 search_by_documentation(self, the_contents, question, min_score=0, how_many_ 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 @@ def check_prime(number): + model_call_counter = {} + AI = AI_() diff --git a/upsonic_on_prem/utils/db/__init__.py b/upsonic_on_prem/utils/db/__init__.py index a0e0442c..852b11f9 100644 --- a/upsonic_on_prem/utils/db/__init__.py +++ b/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")