"""Lisette usage and cost monitoring"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_usage.ipynb.

# %% auto 0
__all__ = ['Usage', 'LisetteUsageLogger']

# %% ../nbs/01_usage.ipynb
from litellm.integrations.custom_logger import CustomLogger
import time
try: 
    from fastlite import *
    from fastlite.core import dataclass
except ImportError: raise ImportError("Please install `fastlite` to use sqlite based lisette usage logging.")

# %% ../nbs/01_usage.ipynb
class Usage: id:int; timestamp:float; model:str; user_id:str; prompt_tokens:int; completion_tokens:int; total_tokens:int; cached_tokens:int; cache_creation_tokens:int; cache_read_tokens:int; web_search_requests:int; response_cost:int

# %% ../nbs/01_usage.ipynb
class LisetteUsageLogger(CustomLogger):
    def __init__(self, db_path): 
        self.db = Database(db_path)
        self.usage = self.db.create(Usage)
    
    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self._log_usage(response_obj, kwargs['response_cost'], start_time, end_time)
    def log_success_event(self, kwargs, response_obj, start_time, end_time):             self._log_usage(response_obj, kwargs['response_cost'], start_time, end_time)
    def _log_usage(self, response_obj, response_cost, start_time, end_time):
        usage = response_obj.usage
        ptd   = usage.prompt_tokens_details
        self.usage.insert(Usage(timestamp=time.time(), model=response_obj.model, user_id=self.user_id_fn(), prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens,
                                    total_tokens=usage.total_tokens, cached_tokens=ptd.cached_tokens if ptd else 0, cache_creation_tokens=usage.cache_creation_input_tokens, 
                                    cache_read_tokens=usage.cache_read_input_tokens, web_search_requests=nested_idx(usage, 'server_tool_use', 'web_search_requests'), response_cost=response_cost))
                  
    def user_id_fn(self): raise NotImplementedError('Please implement `LisetteUsageLogger.user_id_fn` before initializing, e.g using fastcore.patch.')

# %% ../nbs/01_usage.ipynb
@patch
def total_cost(self:Usage, sc=0.01): return self.response_cost + sc * ifnone(self.web_search_requests, 0)
