| from shortGPT.gpt import gpt_utils | |
| from shortGPT.database.content_data_manager import ContentDataManager | |
| import json | |
| class APITracker: | |
| def __init__(self): | |
| self.initiateAPITracking() | |
| def setDataManager(self, contentManager : ContentDataManager): | |
| if(not contentManager): | |
| raise Exception("contentManager is null") | |
| self.datastore = contentManager | |
| def openAIWrapper(self, gptFunc): | |
| def wrapper(*args, **kwargs): | |
| result = gptFunc(*args, **kwargs) | |
| prompt = kwargs.get('prompt') or kwargs.get('conversation') or args[0] | |
| prompt = json.dumps(prompt) | |
| if self.datastore and result: | |
| tokensUsed = gpt_utils.num_tokens_from_messages([prompt, result]) | |
| self.datastore.save('api_openai', tokensUsed, add=True) | |
| return result | |
| return wrapper | |
| def elevenWrapper(self, audioFunc): | |
| def wrapper(*args, **kwargs): | |
| result = audioFunc(*args, **kwargs) | |
| textInput = kwargs.get('text') or args[0] | |
| if self.datastore and result: | |
| self.datastore.save('api_eleven', len(textInput), add=True) | |
| return result | |
| return wrapper | |
| def wrap_turbo(self): | |
| func_name = "gpt3Turbo_completion" | |
| module = __import__("gpt_utils", fromlist=["gpt3Turbo_completion"]) | |
| func = getattr(module, func_name) | |
| wrapped_func = self.openAIWrapper(func) | |
| setattr(module, func_name, wrapped_func) | |
| def wrap_eleven(self): | |
| func_name = "generateVoice" | |
| module = __import__("audio_generation", fromlist=["generateVoice"]) | |
| func = getattr(module, func_name) | |
| wrapped_func = self.elevenWrapper(func) | |
| setattr(module, func_name, wrapped_func) | |
| def initiateAPITracking(self): | |
| self.wrap_turbo() | |
| self.wrap_eleven() | |