import functools

from langchain.callbacks import get_openai_callback


def with_openai_stats(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        with get_openai_callback() as cb:
            result = func(self, *args, **kwargs)
            prompt_token = cb.prompt_tokens
            completion_token = cb.completion_tokens
            total_cost = cb.total_cost
            self.usage = {
                "prompt_tokens": prompt_token,
                "completion_tokens": completion_token,
                "total_cost": total_cost,
            }
        return result

    return wrapper