shortGPT/tracking/api_tracking.py

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()